diff --git a/README.md b/README.md
index fd86de63279ecdd391d7ff9f48da75296a87589a..378184b1faf260afdf6a27748a732018c584001e 100644
--- a/README.md
+++ b/README.md
@@ -1,13 +1,13 @@
 ---
 title: Streaming Code World Model
-emoji: 🐨
-colorFrom: yellow
-colorTo: indigo
+emoji: 🌍
+colorFrom: blue
+colorTo: red
 sdk: gradio
-sdk_version: 6.30.0
-python_version: '3.12'
+sdk_version: 5.50.0
 app_file: app.py
-pinned: false
+short_description: Causal chunk-by-chunk world-model video streaming
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+(prep stage)
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..664b7fe1dd71db6bb5bb6fb180f486d494289b6d
--- /dev/null
+++ b/app.py
@@ -0,0 +1,124 @@
+"""TEMPORARY prep stage: encode the conditioning caches of a few released CWM examples, exactly as
+`python -m rtcwm.eval.open_world prepare` does (Qwen3-VL-32B states of the static caption, Picture 1 latent,
+proxy (Depth + Semantic-ID) latents, silent audio latent). Replaced by the real demo afterwards."""
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import spaces
+import io
+import json
+import shutil
+import tempfile
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+import gradio as gr
+from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import Image
+from safetensors.torch import save_file
+
+H3_REPO, H3_REV = "Comfy-Org/MiniMax-H3", "0543966fbdce5ba05709a8f2031c94bdba629b4a"
+QWEN_REPO, QWEN_REV = "Qwen/Qwen3-VL-32B-Instruct", "0cfaf48183f594c314753d30a4c4974bc75f3ccb"
+EX_REPO, EX_REV = "NTU-yiwen/code-world-model-inference-examples-40", "f8307bf50cb19bed191682734fdf444738160f91"
+EXAMPLES = [3, 6, 8, 29, 31, 36]
+CHUNKS = 24
+CLIP_FRAMES = 17
+
+models = Path(snapshot_download(H3_REPO, revision=H3_REV, allow_patterns=[
+    "text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors", "vae/minimax_h3_video_vae_fp16.safetensors",
+    "vae/minimax_h3_audio_vae_fp32.safetensors"]))
+proc_dir = Path(snapshot_download(QWEN_REPO, revision=QWEN_REV, allow_patterns=[
+    "chat_template.json", "config.json", "generation_config.json", "merges.txt", "preprocessor_config.json",
+    "tokenizer_config.json", "tokenizer.json", "video_preprocessor_config.json", "vocab.json"]))
+src = Path(snapshot_download(EX_REPO, repo_type="dataset", revision=EX_REV, allow_patterns=[
+    "manifest.json", "semantic_id_palette.json"] + [f"{i}/*" for i in EXAMPLES] + [f"{i}/**" for i in EXAMPLES]))
+
+import materialize_hf_examples as mat
+
+REL = Path("/tmp/release")
+for i in EXAMPLES:
+    mat.materialize_example(src, REL, i, require_config=False, frame_limit=(CHUNKS + 1) * CLIP_FRAMES)
+print("materialized", flush=True)
+
+from musubi_tuner.minimax_h3.audio_vae import load_audio_vae
+from musubi_tuner.minimax_h3.text_encoder import load_h3_processor, load_h3_text_encoder
+from musubi_tuner.minimax_h3.video_vae import encode_video_condition, load_video_vae
+
+vae = load_video_vae(models / "vae/minimax_h3_video_vae_fp16.safetensors", device="cpu", dtype=torch.float16).eval().requires_grad_(False)
+vae.to("cuda")
+avae = load_audio_vae(models / "vae/minimax_h3_audio_vae_fp32.safetensors", device="cpu", dtype=torch.float32).eval()
+avae.to("cuda")
+proc = load_h3_processor(str(proc_dir))
+qwen = load_h3_text_encoder(str(models / "text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors"),
+                            processor_path=str(proc_dir), device="cpu", dtype=torch.bfloat16).eval()
+qwen.to("cuda")
+print("models loaded", flush=True)
+
+
+def _prepare_pixels(frames):  # musubi_tuner.minimax_h3_generate_video._prepare_pixels
+    if frames.dtype == torch.uint8:
+        frames = frames.float().div(127.5).sub(1.0)
+    else:
+        frames = frames.float().mul(2.0).sub(1.0)
+    return frames.permute(3, 0, 1, 2).unsqueeze(0).contiguous()
+
+
+@spaces.GPU(duration=600, size="xlarge")
+def prep() -> str:
+    """Encode the example caches and return a zip."""
+    from cwm_h3_inference.cache import _load_anchor
+    from cwm_h3_inference.config import _canonical_caption
+    from cwm_h3_inference.constants import TARGET_HEIGHT, TARGET_WIDTH, WINDOW_AUDIO_SAMPLES
+    from cwm_h3_inference.duv import load_duv_frame
+    from rtcwm.eval.open_world import static_caption
+    from rtcwm.train.preencode import static_text_states
+
+    dev = torch.device("cuda")
+    out = Path(tempfile.mkdtemp())
+    with torch.inference_mode():
+        a = avae.encode(torch.zeros(1, 2, WINDOW_AUDIO_SAMPLES, device=dev)).float().cpu()[0]
+    save_file({"audio": a.contiguous()}, str(out / "silence_audio.safetensors"))
+    log = []
+    for i in EXAMPLES:
+        ex = REL / "examples" / f"example_{i:02d}"
+        cond = ex / "conditions"
+        prompt = (src / str(i) / "prompts" / "000.txt").read_text()
+        k = min(CHUNKS, len(list(cond.glob("*.semantic_id.png"))) // CLIP_FRAMES)
+        duv = np.stack([load_duv_frame(cond, f, for_qwen=False) for f in range(k * CLIP_FRAMES)])
+        duv = np.concatenate([duv, np.repeat(duv[-1:], CLIP_FRAMES, 0)], 0)
+        with torch.inference_mode():
+            ref = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(duv)).to(dev, torch.float16))[0]
+        ref = ref.float().cpu()[:, :5 * k].contiguous()
+        picture1 = _load_anchor(ex / "first_frame.png", size=(TARGET_WIDTH, TARGET_HEIGHT))
+        png = out / f"ex{i:02d}_anchor.png"
+        Image.fromarray(picture1).save(png)
+        with torch.inference_mode():
+            lat = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(picture1).unsqueeze(0)).to(dev, torch.float16))
+        caption = _canonical_caption(static_caption(prompt), f"ow_example_{i:02d}/static_c00")
+        hidden, tags = static_text_states(proc, qwen, picture1, png, cond, 0, caption)
+        sem = np.stack([np.asarray(Image.open(cond / f"{f:06d}.semantic_id.png")) for f in range(k * CLIP_FRAMES)])
+        np.savez_compressed(out / f"ex{i:02d}_semantic.npz", semantic_id=sem.astype(np.uint8))
+        save_file({"duv": ref, "anchor": lat[0].float().cpu().contiguous(),
+                   "hidden": hidden.to("cpu", torch.bfloat16).contiguous(), "tags": tags.to("cpu", torch.int64).contiguous()},
+                  str(out / f"ex{i:02d}.safetensors"), metadata={"chunks": str(k), "caption": caption})
+        log.append(dict(id=i, chunks=k, duv=list(ref.shape), anchor=list(lat.shape), hidden=list(hidden.shape)))
+        print(log[-1], flush=True)
+    (out / "log.json").write_text(json.dumps(log))
+    z = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name
+    with zipfile.ZipFile(z, "w", zipfile.ZIP_STORED) as zf:
+        for p in out.iterdir():
+            zf.write(p, p.name)
+    return z
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("Preparing example caches (temporary).")
+    b = gr.Button("prep")
+    f = gr.File()
+    b.click(prep, None, f, api_name="prep")
+
+demo.launch()
diff --git a/cwm_h3_inference/__init__.py b/cwm_h3_inference/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9b72494ad986cc8510cb4135f6ae9caa6e9e306
--- /dev/null
+++ b/cwm_h3_inference/__init__.py
@@ -0,0 +1,6 @@
+"""MiniMax H3 inference."""
+
+from .config import RunConfig, load_config
+
+__all__ = ["RunConfig", "load_config"]
+__version__ = "0.1.0"
diff --git a/cwm_h3_inference/__main__.py b/cwm_h3_inference/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb53e2f31b2f703ad32ef64b8a41faa3e7d18e08
--- /dev/null
+++ b/cwm_h3_inference/__main__.py
@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())
diff --git a/cwm_h3_inference/cache.py b/cwm_h3_inference/cache.py
new file mode 100644
index 0000000000000000000000000000000000000000..b8b6fd6dcd7d4d3c563866fb704b2ce4823c21b0
--- /dev/null
+++ b/cwm_h3_inference/cache.py
@@ -0,0 +1,321 @@
+"""Prepare the DUV, RGB-anchor, fixed-first-frame, and Qwen caches."""
+
+from __future__ import annotations
+
+import gc
+import hashlib
+from io import BytesIO
+import json
+from pathlib import Path
+from typing import Any
+
+from .config import RunConfig
+from .constants import (
+    ANCHOR_SHORT_EDGE,
+    CACHE_SCHEMA,
+    CANVAS_MULTIPLE,
+    CONDITION_HEIGHT,
+    CONDITION_WIDTH,
+    FPS,
+    QWEN_OFFSETS,
+    REFERENCE_LATENT_SHAPE,
+    SINGLE_FRAME_LATENT_SHAPE,
+    TARGET_HEIGHT,
+    TARGET_WIDTH,
+    VIDEO_LATENT_SHAPE,
+)
+from .duv import load_duv_window
+from .io import atomic_json, atomic_safetensors, load_json
+from .presentation import _system_prompt, build_ref2va_presentation
+
+
+def official_short_edge_size(width: int, height: int) -> tuple[int, int]:
+    if min(width, height) <= 0 or width > 4 * height or height > 4 * width:
+        raise ValueError("anchor aspect ratio must be between 1:4 and 4:1")
+    scale = ANCHOR_SHORT_EDGE / min(width, height)
+
+    def round_axis(value: float) -> int:
+        return max(CANVAS_MULTIPLE, round(value / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
+
+    return round_axis(width * scale), round_axis(height * scale)
+
+
+def _load_anchor(path: Path, *, size: tuple[int, int] | None = None):
+    import numpy as np
+    from PIL import Image
+
+    payload = path.read_bytes()
+    with Image.open(BytesIO(payload)) as image:
+        image = image.convert("RGB")
+        target = official_short_edge_size(*image.size) if size is None else size
+        pixels = np.asarray(image.resize(target, Image.Resampling.LANCZOS), dtype=np.uint8).copy()
+    return pixels
+
+
+def cache_paths(config: RunConfig, window_index: int) -> dict[str, Path]:
+    root = config.cache_dir
+    window = root / f"window_{window_index:03d}"
+    return {
+        "anchor": root / "anchor_reference.safetensors",
+        "input_video": root / "input_video.safetensors",
+        "reference": window / "duv_reference.safetensors",
+        "qwen": window / "qwen.safetensors",
+        "manifest": root / "manifest.json",
+    }
+
+
+def _cache_contract(config: RunConfig) -> dict[str, Any]:
+    return {
+        "video_vae": str(config.models.video_vae),
+        "text_encoder": str(config.models.text_encoder),
+        "processor": str(config.models.processor),
+        "anchor_image": str(config.anchor_image),
+        "anchor_bytes": config.anchor_image.stat().st_size,
+        "condition_root": str(config.condition_root),
+        "system_prompt_w0": _system_prompt(0),
+        "system_prompt_wn": _system_prompt(1),
+        "windows": [
+            {"prompt": window.prompt, "start_frame": window.start_frame}
+            for window in config.windows
+        ],
+    }
+
+
+def _contract_id(value: dict[str, Any]) -> str:
+    payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()
+    return hashlib.sha256(payload).hexdigest()
+
+
+def _begin_manifest(config: RunConfig) -> dict[str, Any]:
+    path = cache_paths(config, 0)["manifest"]
+    contract = _cache_contract(config)
+    contract_id = _contract_id(contract)
+    if path.exists():
+        existing = load_json(path)
+        if existing.get("schema") != CACHE_SCHEMA or existing.get("contract_id") != contract_id:
+            raise ValueError(f"cache directory belongs to another run: {config.cache_dir}")
+        return existing
+    if config.cache_dir.exists() and any(config.cache_dir.iterdir()):
+        raise ValueError(f"non-empty cache directory has no manifest: {config.cache_dir}")
+    manifest = {
+        "schema": CACHE_SCHEMA,
+        "status": "preparing",
+        "contract_id": contract_id,
+        "contract": contract,
+        "window_count": len(config.windows),
+        "output_frames": config.output_frames,
+        "files": [],
+    }
+    atomic_json(path, manifest)
+    return manifest
+
+
+def _require_one_gpu(torch: Any) -> Any:
+    if not torch.cuda.is_available() or torch.cuda.device_count() != 1:
+        raise RuntimeError("cache preparation requires exactly one visible CUDA GPU")
+    torch.cuda.set_device(0)
+    return torch.device("cuda:0")
+
+
+def _prepare_visual_caches(config: RunConfig) -> None:
+    import torch
+    from musubi_tuner.minimax_h3.video_vae import encode_video_condition, load_video_vae
+    from musubi_tuner.minimax_h3_generate_video import _prepare_pixels
+
+    device = _require_one_gpu(torch)
+    vae = load_video_vae(config.models.video_vae, device=device, dtype=torch.float16, disable_mmap=False)
+    vae.eval().requires_grad_(False)
+    common = {"schema": CACHE_SCHEMA}
+    first_paths = cache_paths(config, 0)
+
+    if not first_paths["anchor"].exists():
+        anchor_pixels = _load_anchor(config.anchor_image)
+        frames = torch.from_numpy(anchor_pixels).unsqueeze(0)
+        with torch.inference_mode():
+            latent = encode_video_condition(vae, _prepare_pixels(frames).to(device, torch.float16))[0].cpu()
+        if latent.ndim != 4 or tuple(latent.shape[:2]) != (24, 1):
+            raise RuntimeError(f"unexpected anchor latent shape: {tuple(latent.shape)}")
+        atomic_safetensors(
+            first_paths["anchor"],
+            {"latents_ref_000_image": latent.float().contiguous()},
+            {**common, "kind": "anchor", "source": config.anchor_image.name},
+        )
+        del anchor_pixels, frames, latent
+
+    if not first_paths["input_video"].exists():
+        pixels = _load_anchor(config.anchor_image, size=(TARGET_WIDTH, TARGET_HEIGHT))
+        frames = torch.from_numpy(pixels).unsqueeze(0)
+        prepared = _prepare_pixels(frames).to(device, torch.float16)
+        with torch.inference_mode():
+            single = vae.encode(prepared).float().cpu()
+        if tuple(single.shape) == (1, *SINGLE_FRAME_LATENT_SHAPE):
+            single = single[0]
+        if tuple(single.shape) != SINGLE_FRAME_LATENT_SHAPE:
+            raise RuntimeError(f"unexpected first-frame latent shape: {tuple(single.shape)}")
+        full = torch.zeros(VIDEO_LATENT_SHAPE[1:], dtype=torch.float32)
+        full[:, 0:1] = single
+        atomic_safetensors(
+            first_paths["input_video"],
+            {"input_video_latents": full.contiguous()},
+            {
+                **common,
+                "kind": "fixed_first_frame",
+                "fixed_video_latent_indices": "0",
+                "denoise_mask_semantics": "0_fixed_1_regenerate",
+            },
+        )
+        del pixels, frames, prepared, single, full
+
+    for index, window in enumerate(config.windows):
+        path = cache_paths(config, index)["reference"]
+        if path.exists():
+            continue
+        visual = load_duv_window(config.condition_root, window.start_frame, for_qwen=False)
+        pixels = _prepare_pixels(torch.from_numpy(visual.copy())).to(device, torch.float16)
+        with torch.inference_mode():
+            latent = encode_video_condition(vae, pixels)[0].cpu()
+        if tuple(latent.shape) != REFERENCE_LATENT_SHAPE:
+            raise RuntimeError(f"unexpected DUV reference shape for window {index}: {tuple(latent.shape)}")
+        atomic_safetensors(
+            path,
+            {"latents_ref_001_video": latent.float().contiguous()},
+            {**common, "kind": "duv_reference", "window_index": str(index)},
+        )
+        del visual, pixels, latent
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    del vae
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def _prepare_qwen_caches(config: RunConfig) -> None:
+    import numpy as np
+    import torch
+    from musubi_tuner.minimax_h3.media import H3AudioSource, H3Record, H3Reference
+    from musubi_tuner.minimax_h3.text_encoder import (
+        H3TextVisual,
+        build_presentation,
+        encode_h3_presentation,
+        load_h3_processor,
+        load_h3_text_encoder,
+    )
+
+    pending = [index for index in range(len(config.windows)) if not cache_paths(config, index)["qwen"].exists()]
+    if not pending:
+        return
+    device = _require_one_gpu(torch)
+    processor = load_h3_processor(str(config.models.processor))
+    model = load_h3_text_encoder(
+        str(config.models.text_encoder),
+        processor_path=str(config.models.processor),
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    model.eval().requires_grad_(False)
+    anchor_np = _load_anchor(config.anchor_image)
+    anchor = torch.from_numpy(np.ascontiguousarray(anchor_np)).unsqueeze(0)
+    for index in pending:
+        window = config.windows[index]
+        duv_np = load_duv_window(config.condition_root, window.start_frame, for_qwen=True)
+        duv = torch.from_numpy(np.ascontiguousarray(duv_np))
+        image_key = config.anchor_image
+        video_key = config.condition_root / f"{window.start_frame:06d}.semantic_id.png"
+        placeholder = config.condition_root / ".h3-presentation-placeholder"
+        record = H3Record(
+            video_path=placeholder,
+            caption=window.prompt,
+            target_audio=H3AudioSource(path=placeholder, embedded=False),
+            references=(
+                H3Reference(type="image", path=image_key),
+                H3Reference(type="video", path=video_key, duration_seconds=124 / FPS),
+            ),
+            jsonl_line=0,
+        )
+        plain = build_presentation(
+            record,
+            "ref2va",
+            {
+                image_key: H3TextVisual(anchor),
+                video_key: H3TextVisual(duv, tuple(offset / FPS for offset in QWEN_OFFSETS)),
+            },
+        )
+        presentation = build_ref2va_presentation(processor, plain, index, window.prompt)
+        with torch.inference_mode():
+            hidden, tags = encode_h3_presentation(processor, model, presentation)
+        hidden = hidden.to(device="cpu", dtype=torch.bfloat16).contiguous()
+        tags = tags.to(device="cpu", dtype=torch.int64).contiguous()
+        if hidden.ndim != 2 or hidden.shape[1] != 5120 or tags.shape != (hidden.shape[0],):
+            raise RuntimeError(f"unexpected Qwen output shape for window {index}")
+        atomic_safetensors(
+            cache_paths(config, index)["qwen"],
+            {"varlen_mmh3_hidden_states_bfloat16": hidden, "varlen_mmh3_token_tags_int64": tags},
+            {
+                "schema": CACHE_SCHEMA,
+                "kind": "qwen",
+                "window_index": str(index),
+                "system_prompt_role": "w0" if index == 0 else "wn",
+                "presentation_order": "picture1_anchor_then_video1_duv",
+            },
+        )
+        del duv_np, duv, plain, presentation, hidden, tags
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    del anchor_np, anchor, processor, model
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def prepare_caches(config: RunConfig) -> dict[str, Any]:
+    config.validate_files()
+    initial = _begin_manifest(config)
+    if initial.get("status") == "complete":
+        return validate_cache(config)
+    _prepare_visual_caches(config)
+    _prepare_qwen_caches(config)
+    files = []
+    for index in range(len(config.windows)):
+        paths = cache_paths(config, index)
+        files.extend([paths["reference"], paths["qwen"]])
+    files.extend([cache_paths(config, 0)["anchor"], cache_paths(config, 0)["input_video"]])
+    missing = [str(path) for path in files if not path.is_file()]
+    if missing:
+        raise RuntimeError(f"cache preparation incomplete: {missing}")
+    manifest = {
+        "schema": CACHE_SCHEMA,
+        "status": "complete",
+        "contract_id": initial["contract_id"],
+        "contract": initial["contract"],
+        "window_count": len(config.windows),
+        "output_frames": config.output_frames,
+        "files": [str(path) for path in files],
+    }
+    atomic_json(cache_paths(config, 0)["manifest"], manifest)
+    return manifest
+
+
+def validate_cache(config: RunConfig) -> dict[str, Any]:
+    manifest = load_json(cache_paths(config, 0)["manifest"])
+    expected_id = _contract_id(_cache_contract(config))
+    if (
+        manifest.get("schema") != CACHE_SCHEMA
+        or manifest.get("status") != "complete"
+        or manifest.get("contract_id") != expected_id
+    ):
+        raise ValueError("cache manifest is incompatible with this release")
+    if manifest.get("window_count") != len(config.windows):
+        raise ValueError("cache window count does not match config")
+    expected_files: list[Path] = []
+    for index in range(len(config.windows)):
+        paths = cache_paths(config, index)
+        expected_files.extend([paths["reference"], paths["qwen"]])
+    expected_files.extend([cache_paths(config, 0)["anchor"], cache_paths(config, 0)["input_video"]])
+    files = manifest.get("files")
+    if files != [str(path) for path in expected_files]:
+        raise ValueError("cache manifest file set does not match config")
+    missing = [str(path) for path in expected_files if not path.is_file()]
+    if missing:
+        raise ValueError(f"cache files are missing: {missing}")
+    return manifest
diff --git a/cwm_h3_inference/cli.py b/cwm_h3_inference/cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..60385aba26af7b1842f2575efbd5e834e8955b92
--- /dev/null
+++ b/cwm_h3_inference/cli.py
@@ -0,0 +1,53 @@
+"""Command-line entry point."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from typing import Any
+
+from .cache import prepare_caches, validate_cache
+from .config import load_config
+from .engine import generate
+
+
+def _emit(value: Any) -> None:
+    print(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True))
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="cwm-h3-inference",
+        description="MiniMax H3 Ref2VA inference with a fixed first frame and Retake34 continuation.",
+        allow_abbrev=False,
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+    for name in ("validate", "prepare", "generate"):
+        command = subparsers.add_parser(name, allow_abbrev=False)
+        command.add_argument("--config", required=True, help="JSON run configuration")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    config = load_config(args.config, check_files=True)
+    if args.command == "validate":
+        _emit(
+            {
+                "status": "pass",
+                "window_count": len(config.windows),
+                "output_frames": config.output_frames,
+            }
+        )
+        return 0
+    if args.command == "prepare":
+        _emit(prepare_caches(config))
+        return 0
+    if args.command == "generate":
+        _emit(generate(config))
+        return 0
+    raise AssertionError(f"unhandled command: {args.command}")
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/cwm_h3_inference/config.py b/cwm_h3_inference/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f6df2f2ad9c50292577e91cc72d95ad7c83d729
--- /dev/null
+++ b/cwm_h3_inference/config.py
@@ -0,0 +1,174 @@
+"""Structured, path-parameterized configuration for one inference run."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+from typing import Any
+
+from .constants import DEFAULT_STEPS, EXPECTED_LORA_MODULES, STRIDE_FRAMES, WINDOW_FRAMES
+
+
+ALLOWED_ROOT_KEYS = {
+    "models",
+    "anchor_image",
+    "condition_root",
+    "cache_dir",
+    "output",
+    "windows",
+    "seed",
+    "steps",
+    "attention_mode",
+    "blocks_to_swap",
+    "expected_lora_modules",
+}
+ALLOWED_MODEL_KEYS = {"dit", "video_vae", "audio_vae", "text_encoder", "processor", "lora"}
+
+
+def require(condition: bool, message: str) -> None:
+    if not condition:
+        raise ValueError(message)
+
+
+def _path(value: Any, label: str, *, base: Path) -> Path:
+    require(isinstance(value, str) and bool(value.strip()), f"{label} must be a non-empty path")
+    path = Path(value)
+    return path if path.is_absolute() else (base / path).resolve()
+
+
+def _canonical_caption(value: Any, label: str) -> str:
+    require(isinstance(value, str) and bool(value.strip()), f"{label} must be non-empty")
+    # The released Qwen caches were encoded from UTF-8 caption bytes whose
+    # internal line endings are CRLF. JSON editors and Path.read_text() commonly
+    # normalize those bytes to LF, which changes Qwen token IDs. Canonicalize the
+    # user caption here so both representations consume the released text.
+    normalized = value.strip().replace("\r\n", "\n").replace("\r", "\n")
+    return normalized.replace("\n", "\r\n")
+
+
+@dataclass(frozen=True)
+class ModelPaths:
+    dit: Path
+    video_vae: Path
+    audio_vae: Path
+    text_encoder: Path
+    processor: Path
+    lora: Path
+
+
+@dataclass(frozen=True)
+class Window:
+    prompt: str
+    start_frame: int
+
+
+@dataclass(frozen=True)
+class RunConfig:
+    source: Path
+    models: ModelPaths
+    anchor_image: Path
+    condition_root: Path
+    cache_dir: Path
+    output: Path
+    windows: tuple[Window, ...]
+    seed: int
+    steps: int
+    attention_mode: str
+    blocks_to_swap: int
+    expected_lora_modules: int
+
+    @property
+    def output_frames(self) -> int:
+        from .constants import STRIDE_FRAMES
+
+        return WINDOW_FRAMES + (len(self.windows) - 1) * STRIDE_FRAMES
+
+    def validate_files(self) -> None:
+        for label, path in (
+            ("dit", self.models.dit),
+            ("video_vae", self.models.video_vae),
+            ("audio_vae", self.models.audio_vae),
+            ("text_encoder", self.models.text_encoder),
+            ("processor", self.models.processor),
+            ("lora", self.models.lora),
+            ("anchor_image", self.anchor_image),
+            ("condition_root", self.condition_root),
+        ):
+            require(path.exists(), f"{label} does not exist: {path}")
+        from .presentation import _system_prompt
+
+        _system_prompt(0)
+        _system_prompt(1)
+
+
+def load_config(path: str | Path, *, check_files: bool = False) -> RunConfig:
+    source = Path(path).resolve()
+    raw = json.loads(source.read_text(encoding="utf-8"))
+    require(isinstance(raw, dict), "config root must be an object")
+    require(not (set(raw) - ALLOWED_ROOT_KEYS), f"unknown config keys: {sorted(set(raw) - ALLOWED_ROOT_KEYS)}")
+    base = source.parent
+
+    model_raw = raw.get("models")
+    require(isinstance(model_raw, dict), "models must be an object")
+    require(
+        set(model_raw) == ALLOWED_MODEL_KEYS,
+        f"models must contain exactly: {sorted(ALLOWED_MODEL_KEYS)}",
+    )
+    models = ModelPaths(
+        dit=_path(model_raw.get("dit"), "models.dit", base=base),
+        video_vae=_path(model_raw.get("video_vae"), "models.video_vae", base=base),
+        audio_vae=_path(model_raw.get("audio_vae"), "models.audio_vae", base=base),
+        text_encoder=_path(model_raw.get("text_encoder"), "models.text_encoder", base=base),
+        processor=_path(model_raw.get("processor"), "models.processor", base=base),
+        lora=_path(model_raw.get("lora"), "models.lora", base=base),
+    )
+
+    window_raw = raw.get("windows")
+    require(isinstance(window_raw, list) and bool(window_raw), "windows must be a non-empty list")
+    windows: list[Window] = []
+    for index, item in enumerate(window_raw):
+        require(isinstance(item, dict), f"windows[{index}] must be an object")
+        prompt = item.get("prompt")
+        start = item.get("start_frame")
+        prompt = _canonical_caption(prompt, f"windows[{index}].prompt")
+        require(type(start) is int and start >= 0, f"windows[{index}].start_frame must be non-negative")
+        require(set(item) == {"prompt", "start_frame"}, f"windows[{index}] has unknown or missing keys")
+        if windows:
+            require(
+                start == windows[-1].start_frame + STRIDE_FRAMES,
+                f"windows[{index}].start_frame must advance by {STRIDE_FRAMES}",
+            )
+        windows.append(Window(prompt=prompt, start_frame=start))
+
+    seed = raw.get("seed", 0)
+    steps = raw.get("steps", DEFAULT_STEPS)
+    attention_mode = raw.get("attention_mode", "flash3")
+    blocks_to_swap = raw.get("blocks_to_swap", 48)
+    expected_lora_modules = raw.get("expected_lora_modules", EXPECTED_LORA_MODULES)
+    require(type(seed) is int and 0 <= seed < 2**63, "seed must be an integer in [0, 2**63)")
+    require(type(steps) is int and steps == DEFAULT_STEPS, "this release fixes explicit Euler inference to 20 steps")
+    require(attention_mode == "flash3", "this release fixes attention_mode to flash3")
+    require(type(blocks_to_swap) is int and 0 <= blocks_to_swap <= 48, "blocks_to_swap must be in [0,48]")
+    require(
+        type(expected_lora_modules) is int and expected_lora_modules == EXPECTED_LORA_MODULES,
+        f"expected_lora_modules must be {EXPECTED_LORA_MODULES}",
+    )
+
+    config = RunConfig(
+        source=source,
+        models=models,
+        anchor_image=_path(raw.get("anchor_image"), "anchor_image", base=base),
+        condition_root=_path(raw.get("condition_root"), "condition_root", base=base),
+        cache_dir=_path(raw.get("cache_dir", "cache"), "cache_dir", base=base),
+        output=_path(raw.get("output", "output.mp4"), "output", base=base),
+        windows=tuple(windows),
+        seed=seed,
+        steps=steps,
+        attention_mode=attention_mode,
+        blocks_to_swap=blocks_to_swap,
+        expected_lora_modules=expected_lora_modules,
+    )
+    if check_files:
+        config.validate_files()
+    return config
diff --git a/cwm_h3_inference/constants.py b/cwm_h3_inference/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..b76f5c87c820ac4f7caa6b51c864e3241b0121f1
--- /dev/null
+++ b/cwm_h3_inference/constants.py
@@ -0,0 +1,43 @@
+"""Frozen geometry for the released Ref2VA inference path."""
+
+FPS = 24
+SAMPLE_RATE = 32_000
+WINDOW_FRAMES = 124
+OVERLAP_FRAMES = 34
+STRIDE_FRAMES = WINDOW_FRAMES - OVERLAP_FRAMES
+
+CONDITION_WIDTH = 336
+CONDITION_HEIGHT = 192
+TARGET_WIDTH = 1_344
+TARGET_HEIGHT = 768
+ANCHOR_SHORT_EDGE = 2_048
+CANVAS_MULTIPLE = 32
+
+VIDEO_LATENT_SHAPE = (1, 24, 37, 48, 84)
+AUDIO_LATENT_SHAPE = (1, 32, 2, 207)
+REFERENCE_LATENT_SHAPE = (24, 37, 12, 21)
+REFERENCE_RUNTIME_SHAPE = (1, 24, 37, 12, 22)
+SINGLE_FRAME_LATENT_SHAPE = (24, 1, 48, 84)
+
+VIDEO_PREFIX_LATENTS = 10
+AUDIO_PREFIX_LATENTS = 57
+WINDOW_AUDIO_SAMPLES = round(WINDOW_FRAMES * SAMPLE_RATE / FPS)
+OVERLAP_AUDIO_SAMPLES = round(OVERLAP_FRAMES * SAMPLE_RATE / FPS)
+QWEN_OFFSETS = tuple(range(0, WINDOW_FRAMES, 12))
+
+DEPTH_NEAR_METRES = 0.3
+DEPTH_FAR_METRES = 256.0
+DEPTH_VALID_EPSILON_METRES = 1.0e-3
+SEMANTIC_U = (32, 96, 160, 224)
+SEMANTIC_V = (43, 128, 213)
+
+DEFAULT_STEPS = 20
+VIDEO_SHIFT = 12.0
+AUDIO_SHIFT = 3.0
+VISUAL_CONDITION_CLEAN = 0.999
+AUDIO_CONDITION_CLEAN = 1.0
+EXPECTED_LORA_MODULES = 200
+
+CACHE_SCHEMA = "cwm_h3_inference_cache_v3"
+RECOVERY_SCHEMA = "cwm_h3_inference_recovery_v3"
+RESULT_SCHEMA = "cwm_h3_inference_result_v3"
diff --git a/cwm_h3_inference/duv.py b/cwm_h3_inference/duv.py
new file mode 100644
index 0000000000000000000000000000000000000000..5bfefe0b1f7028579dc43d8ef62a6a8fbd410384
--- /dev/null
+++ b/cwm_h3_inference/duv.py
@@ -0,0 +1,72 @@
+"""Depth + Semantic-ID to three-channel DUV conditioning."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+from .constants import (
+    CONDITION_HEIGHT,
+    CONDITION_WIDTH,
+    DEPTH_FAR_METRES,
+    DEPTH_NEAR_METRES,
+    DEPTH_VALID_EPSILON_METRES,
+    QWEN_OFFSETS,
+    SEMANTIC_U,
+    SEMANTIC_V,
+    WINDOW_FRAMES,
+)
+
+
+def frame_paths(root: Path, ordinal: int) -> tuple[Path, Path]:
+    return root / f"{ordinal:06d}.depth.f32", root / f"{ordinal:06d}.semantic_id.png"
+
+
+def load_duv_frame(root: Path, ordinal: int, *, for_qwen: bool):
+    import numpy as np
+    from PIL import Image
+
+    depth_path, semantic_path = frame_paths(root, ordinal)
+    payload = depth_path.read_bytes()
+    expected_bytes = CONDITION_WIDTH * CONDITION_HEIGHT * 4
+    if len(payload) != expected_bytes:
+        raise ValueError(f"raw depth byte count must be {expected_bytes}: {depth_path}")
+    metric = np.frombuffer(payload, dtype="<f4")
+    if not bool(np.all(np.isfinite(metric))) or bool(np.any(metric < np.float32(0.0))):
+        raise ValueError(f"raw depth must contain finite non-negative metres: {depth_path}")
+    valid = metric > np.float32(DEPTH_VALID_EPSILON_METRES)
+    codes = np.zeros(metric.shape, dtype=np.uint16)
+    if bool(np.any(valid)):
+        clipped = np.clip(metric[valid].astype(np.float64), DEPTH_NEAR_METRES, DEPTH_FAR_METRES)
+        normalized = (math.log(DEPTH_FAR_METRES) - np.log(clipped)) / (
+            math.log(DEPTH_FAR_METRES) - math.log(DEPTH_NEAR_METRES)
+        )
+        codes[valid] = np.floor(np.clip(normalized, 0.0, 1.0) * 65535.0 + 0.5).astype(np.uint16)
+    depth = codes.reshape(CONDITION_HEIGHT, CONDITION_WIDTH)
+
+    with Image.open(semantic_path) as image:
+        if image.mode != "L" or image.size != (CONDITION_WIDTH, CONDITION_HEIGHT):
+            raise ValueError(f"semantic ID must be L/{CONDITION_WIDTH}x{CONDITION_HEIGHT}: {semantic_path}")
+        semantic = np.asarray(image, dtype=np.uint8).copy()
+    if int(semantic.max()) > 11:
+        raise ValueError(f"semantic ID must be in [0,11]: {semantic_path}")
+    u = np.asarray(SEMANTIC_U, dtype=np.uint8)[semantic % 4]
+    v = np.asarray(SEMANTIC_V, dtype=np.uint8)[semantic // 4]
+    if for_qwen:
+        depth8 = np.floor(depth.astype(np.float64) / 257.0 + 0.5).astype(np.uint8)
+        return np.stack((depth8, u, v), axis=-1)
+    return np.stack(
+        (
+            depth.astype(np.float32) / np.float32(65535.0),
+            u.astype(np.float32) / np.float32(255.0),
+            v.astype(np.float32) / np.float32(255.0),
+        ),
+        axis=-1,
+    )
+
+
+def load_duv_window(root: Path, start_frame: int, *, for_qwen: bool):
+    import numpy as np
+
+    offsets = QWEN_OFFSETS if for_qwen else range(WINDOW_FRAMES)
+    return np.stack([load_duv_frame(root, start_frame + offset, for_qwen=for_qwen) for offset in offsets])
diff --git a/cwm_h3_inference/engine.py b/cwm_h3_inference/engine.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ed4a1cfa3d8535e9c4aac6b5ea2710d17de66ab
--- /dev/null
+++ b/cwm_h3_inference/engine.py
@@ -0,0 +1,340 @@
+"""Single-GPU Ref2VA inference with a fixed first frame and Retake34 continuation."""
+
+from __future__ import annotations
+
+import gc
+import os
+from pathlib import Path
+import time
+from typing import Any
+import uuid
+
+from .cache import cache_paths, validate_cache
+from .config import RunConfig
+from .constants import (
+    AUDIO_CONDITION_CLEAN,
+    AUDIO_LATENT_SHAPE,
+    AUDIO_PREFIX_LATENTS,
+    AUDIO_SHIFT,
+    FPS,
+    OVERLAP_AUDIO_SAMPLES,
+    OVERLAP_FRAMES,
+    RECOVERY_SCHEMA,
+    REFERENCE_LATENT_SHAPE,
+    REFERENCE_RUNTIME_SHAPE,
+    RESULT_SCHEMA,
+    STRIDE_FRAMES,
+    TARGET_HEIGHT,
+    TARGET_WIDTH,
+    VIDEO_LATENT_SHAPE,
+    VIDEO_PREFIX_LATENTS,
+    VIDEO_SHIFT,
+    VISUAL_CONDITION_CLEAN,
+    WINDOW_AUDIO_SAMPLES,
+    WINDOW_FRAMES,
+)
+from .io import atomic_json
+from .runtime import build_runtime_batch, strict_merge_lora
+
+
+def _require_one_gpu(torch: Any) -> Any:
+    if not torch.cuda.is_available() or torch.cuda.device_count() != 1:
+        raise RuntimeError("generation requires exactly one visible CUDA GPU")
+    torch.cuda.set_device(0)
+    return torch.device("cuda:0")
+
+
+def _unique_key(handle: Any, *, exact: str | None = None, prefix: str | None = None) -> str:
+    keys = list(handle.keys())
+    selected = [key for key in keys if key == exact] if exact else [key for key in keys if key.startswith(str(prefix))]
+    if len(selected) != 1:
+        raise ValueError(f"expected one tensor key, got {keys}")
+    return selected[0]
+
+
+class Ref2VAEngine:
+    def __init__(self, config: RunConfig) -> None:
+        import torch
+        from musubi_tuner.minimax_h3.audio_vae import load_audio_vae
+        from musubi_tuner.minimax_h3.model import load_h3_transformer
+        from musubi_tuner.minimax_h3.video_vae import load_video_vae
+        from musubi_tuner.modules.custom_offloading_utils import BlockSwapConfig
+
+        self.config = config
+        self.torch = torch
+        self.device = _require_one_gpu(torch)
+        transformer = load_h3_transformer(
+            config.models.dit,
+            device="cpu",
+            dtype=torch.bfloat16,
+            attn_mode=config.attention_mode,
+            split_attn=False,
+            disable_mmap=False,
+        )
+        self.lora_gate = strict_merge_lora(
+            transformer, config.models.lora, config.expected_lora_modules
+        )
+        swap = BlockSwapConfig(device=self.device, supports_backward=False, use_pinned_memory=True)
+        transformer.enable_block_swap(config.blocks_to_swap, swap)
+        transformer.move_to_device_except_swap_blocks(self.device)
+        transformer.prepare_block_swap_before_forward()
+        transformer.switch_block_swap_for_inference()
+        transformer.eval().requires_grad_(False)
+        self.transformer = transformer
+        self.video_vae = load_video_vae(
+            config.models.video_vae, device=self.device, dtype=torch.float16, disable_mmap=False
+        ).eval().requires_grad_(False)
+        self.audio_vae = load_audio_vae(
+            config.models.audio_vae, device=self.device, dtype=torch.float32, disable_mmap=False
+        ).eval().requires_grad_(False)
+
+    def _load_window_cache(self, index: int):
+        import torch
+        from safetensors import safe_open
+
+        paths = cache_paths(self.config, index)
+        with safe_open(str(paths["anchor"]), framework="pt", device="cpu") as handle:
+            anchor = handle.get_tensor(_unique_key(handle, prefix="latents_ref_"))
+        with safe_open(str(paths["reference"]), framework="pt", device="cpu") as handle:
+            reference = handle.get_tensor(_unique_key(handle, prefix="latents_ref_"))
+        with safe_open(str(paths["qwen"]), framework="pt", device="cpu") as handle:
+            hidden = handle.get_tensor(_unique_key(handle, exact="varlen_mmh3_hidden_states_bfloat16"))
+            tags = handle.get_tensor(_unique_key(handle, exact="varlen_mmh3_token_tags_int64"))
+        if tuple(reference.shape) != REFERENCE_LATENT_SHAPE or reference.dtype != torch.float32:
+            raise ValueError(f"window {index} DUV reference cache shape/dtype changed")
+        reference = torch.nn.functional.pad(reference.unsqueeze(0), (0, 1, 0, 0), value=0.0).contiguous()
+        if tuple(reference.shape) != REFERENCE_RUNTIME_SHAPE:
+            raise ValueError(f"window {index} padded DUV reference shape changed")
+        if anchor.ndim != 4 or tuple(anchor.shape[:2]) != (24, 1) or anchor.dtype != torch.float32:
+            raise ValueError("anchor cache shape/dtype changed")
+        if hidden.ndim != 2 or hidden.shape[1] != 5120 or hidden.dtype != torch.bfloat16:
+            raise ValueError(f"window {index} Qwen hidden states changed")
+        if tags.shape != (hidden.shape[0],) or tags.dtype != torch.int64:
+            raise ValueError(f"window {index} Qwen token tags changed")
+        if not all(torch.isfinite(value).all().item() for value in (anchor, reference, hidden)):
+            raise ValueError(f"window {index} cache contains non-finite values")
+        return anchor.unsqueeze(0), reference, hidden.unsqueeze(0), tags.unsqueeze(0)
+
+    def _load_input_video(self):
+        import torch
+        from safetensors import safe_open
+
+        path = cache_paths(self.config, 0)["input_video"]
+        with safe_open(str(path), framework="pt", device="cpu") as handle:
+            latent = handle.get_tensor(_unique_key(handle, exact="input_video_latents"))
+        if tuple(latent.shape) != VIDEO_LATENT_SHAPE[1:] or latent.dtype != torch.float32:
+            raise ValueError("fixed-first-frame latent shape/dtype changed")
+        if torch.count_nonzero(latent[:, 1:]).item() != 0:
+            raise ValueError("fixed-first-frame latent must contain data only at temporal index 0")
+        latent = latent.unsqueeze(0).to(self.device)
+        mask = torch.ones_like(latent, dtype=torch.bool)
+        mask[:, :, 0] = False
+        return latent, mask
+
+    def _recovery_path(self, index: int) -> Path:
+        return self.config.cache_dir / "recovery" / f"window_{index:03d}.safetensors"
+
+    def _load_recovery(self, index: int):
+        import torch
+        from safetensors import safe_open
+
+        path = self._recovery_path(index)
+        with safe_open(str(path), framework="pt", device="cpu") as handle:
+            metadata = handle.metadata() or {}
+            if metadata.get("schema") != RECOVERY_SCHEMA or metadata.get("window_index") != str(index):
+                raise ValueError(f"recovery identity changed: {path}")
+            video = handle.get_tensor("video")
+            audio = handle.get_tensor("audio")
+        if tuple(video.shape) != VIDEO_LATENT_SHAPE[1:] or tuple(audio.shape) != AUDIO_LATENT_SHAPE[1:]:
+            raise ValueError(f"recovery shape changed: {path}")
+        if video.dtype != torch.float32 or audio.dtype != torch.float32:
+            raise ValueError(f"recovery dtype changed: {path}")
+        if not torch.isfinite(video).all().item() or not torch.isfinite(audio).all().item():
+            raise ValueError(f"recovery contains non-finite values: {path}")
+        return video, audio
+
+    def _retake_inputs(self, previous: tuple[Any, Any]):
+        torch = self.torch
+        previous_video, previous_audio = previous
+        torch.cuda.empty_cache()
+        with torch.inference_mode():
+            decoded_video = self.video_vae.decode(
+                previous_video.unsqueeze(0).to(self.device, torch.float16)
+            )[:, :, :WINDOW_FRAMES]
+            if tuple(decoded_video.shape) != (1, 3, WINDOW_FRAMES, TARGET_HEIGHT, TARGET_WIDTH):
+                raise RuntimeError("Retake34 decoded video shape changed")
+            tail = decoded_video[:, :, -OVERLAP_FRAMES:]
+            padded = torch.cat((tail, tail[:, :, -1:].repeat(1, 1, STRIDE_FRAMES, 1, 1)), dim=2)
+            clean_video = self.video_vae.encode(padded.to(torch.float16)).float()
+            decoded_audio = self.audio_vae.decode(
+                previous_audio.unsqueeze(0).to(self.device, torch.float32)
+            )[:, :, :WINDOW_AUDIO_SAMPLES]
+            if tuple(decoded_audio.shape) != (1, 2, WINDOW_AUDIO_SAMPLES):
+                raise RuntimeError("Retake34 decoded audio shape changed")
+            prefix = decoded_audio[:, :, -OVERLAP_AUDIO_SAMPLES:]
+            prefix_audio = self.audio_vae.encode(prefix.to(torch.float32)).float()
+            if prefix_audio.shape[-1] != AUDIO_PREFIX_LATENTS:
+                raise RuntimeError("Retake34 audio prefix shape changed")
+            clean_audio = torch.cat(
+                (prefix_audio, prefix_audio[..., -1:].repeat(1, 1, 1, AUDIO_LATENT_SHAPE[-1] - AUDIO_PREFIX_LATENTS)),
+                dim=-1,
+            )
+        if tuple(clean_video.shape) != VIDEO_LATENT_SHAPE or tuple(clean_audio.shape) != AUDIO_LATENT_SHAPE:
+            raise RuntimeError("Retake34 clean latent shape changed")
+        video_mask = torch.ones_like(clean_video, dtype=torch.bool)
+        video_mask[:, :, :VIDEO_PREFIX_LATENTS] = False
+        audio_mask = torch.ones_like(clean_audio, dtype=torch.bool)
+        audio_mask[..., :AUDIO_PREFIX_LATENTS] = False
+        return clean_video, video_mask, clean_audio, audio_mask
+
+    def sample_window(self, index: int) -> Path:
+        import torch
+        from musubi_tuner.minimax_h3.sampling import initialize_target_latents, sample_joint_av
+        from safetensors.torch import save_file
+
+        output = self._recovery_path(index)
+        if output.exists():
+            self._load_recovery(index)
+            return output
+        anchor, reference, hidden, tags = self._load_window_cache(index)
+        runtime = build_runtime_batch(anchor, reference, hidden, tags)
+        window_seed = self.config.seed + index
+        initial_video, initial_audio = initialize_target_latents(
+            video_shape=VIDEO_LATENT_SHAPE,
+            audio_shape=AUDIO_LATENT_SHAPE,
+            seed=window_seed,
+            device=self.device,
+            video_dtype=torch.float32,
+            audio_dtype=torch.float32,
+        )
+        ref = reference.to(self.device, torch.float32)
+        ref_generator = torch.Generator(device="cpu").manual_seed(self.config.seed + 1000 + index)
+        ref_noise = torch.randn(tuple(ref.shape), generator=ref_generator, dtype=torch.float32).to(self.device)
+        augmented_ref = VISUAL_CONDITION_CLEAN * ref + (1.0 - VISUAL_CONDITION_CLEAN) * ref_noise
+        anchor_ref = anchor.to(self.device, torch.float32)
+        anchor_generator = torch.Generator(device="cpu").manual_seed(self.config.seed + 3000 + index)
+        anchor_noise = torch.randn(tuple(anchor_ref.shape), generator=anchor_generator, dtype=torch.float32).to(self.device)
+        augmented_anchor = VISUAL_CONDITION_CLEAN * anchor_ref + (1.0 - VISUAL_CONDITION_CLEAN) * anchor_noise
+        retake = self._load_input_video() + (None, None) if index == 0 else self._retake_inputs(self._load_recovery(index - 1))
+        with torch.inference_mode():
+            sampled = sample_joint_av(
+                self.transformer,
+                layout=runtime.layout,
+                text_hidden_states=runtime.text_hidden_states.to(self.device, torch.bfloat16),
+                text_token_tags=runtime.text_token_tags.to(self.device),
+                initial_video=initial_video,
+                initial_audio=initial_audio,
+                steps=self.config.steps,
+                video_shift=VIDEO_SHIFT,
+                audio_shift=AUDIO_SHIFT,
+                visual_condition_latents=(augmented_anchor, augmented_ref),
+                audio_condition_latents=(),
+                visual_condition_clean=VISUAL_CONDITION_CLEAN,
+                audio_condition_clean=AUDIO_CONDITION_CLEAN,
+                input_video_latents=retake[0],
+                video_denoise_mask=retake[1],
+                input_audio_latents=retake[2],
+                audio_denoise_mask=retake[3],
+            )
+        video = sampled.video[0].float().cpu().contiguous()
+        audio = sampled.audio[0].float().cpu().contiguous()
+        if not torch.isfinite(video).all().item() or not torch.isfinite(audio).all().item():
+            raise RuntimeError(f"window {index} produced non-finite latents")
+        output.parent.mkdir(parents=True, exist_ok=True)
+        temporary = output.with_name(f".{output.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+        try:
+            save_file(
+                {"video": video, "audio": audio},
+                str(temporary),
+                metadata={
+                    "schema": RECOVERY_SCHEMA,
+                    "window_index": str(index),
+                    "seed": str(window_seed),
+                    "steps": str(self.config.steps),
+                },
+            )
+            os.replace(temporary, output)
+        finally:
+            temporary.unlink(missing_ok=True)
+            del sampled, video, audio, anchor, reference, hidden, tags, runtime
+            del initial_video, initial_audio, ref, ref_noise, augmented_ref
+            del anchor_ref, anchor_noise, augmented_anchor, retake
+            gc.collect()
+            torch.cuda.empty_cache()
+        return output
+
+    def write_video(self) -> dict[str, Any]:
+        import av
+        import torch
+
+        output = self.config.output
+        if output.exists():
+            raise FileExistsError(f"refusing to overwrite output: {output}")
+        output.parent.mkdir(parents=True, exist_ok=True)
+        temporary = output.with_name(f".{output.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp.mp4")
+        frame_count = 0
+        try:
+            with av.open(str(temporary), mode="w") as container:
+                stream = container.add_stream("libx264", rate=FPS)
+                stream.width = TARGET_WIDTH
+                stream.height = TARGET_HEIGHT
+                stream.pix_fmt = "yuv420p"
+                for index in range(len(self.config.windows)):
+                    video, _audio = self._load_recovery(index)
+                    torch.cuda.empty_cache()
+                    with torch.inference_mode():
+                        decoded = self.video_vae.decode(
+                            video.unsqueeze(0).to(self.device, torch.float16)
+                        )[0, :, :WINDOW_FRAMES].cpu()
+                    pixels = (
+                        ((decoded.float().clamp(-1, 1) + 1.0) * 127.5)
+                        .round()
+                        .to(torch.uint8)
+                        .permute(1, 2, 3, 0)
+                        .contiguous()
+                    )
+                    if index:
+                        pixels = pixels[OVERLAP_FRAMES:]
+                    for value in pixels:
+                        frame = av.VideoFrame.from_ndarray(value.numpy(), format="rgb24")
+                        for packet in stream.encode(frame):
+                            container.mux(packet)
+                        frame_count += 1
+                    del video, decoded, pixels
+                    gc.collect()
+                    torch.cuda.empty_cache()
+                for packet in stream.encode():
+                    container.mux(packet)
+            if frame_count != self.config.output_frames:
+                raise RuntimeError(f"stitched frame count mismatch: {frame_count}")
+            os.replace(temporary, output)
+        finally:
+            temporary.unlink(missing_ok=True)
+        receipt = {
+            "schema": RESULT_SCHEMA,
+            "solver": "explicit_euler",
+            "steps": self.config.steps,
+            "seed": self.config.seed,
+            "window_frames": WINDOW_FRAMES,
+            "overlap_frames": OVERLAP_FRAMES,
+            "output_frames": frame_count,
+            "fps": FPS,
+            "resolution": [TARGET_WIDTH, TARGET_HEIGHT],
+            "audio": False,
+            "output": str(output),
+            "bytes": output.stat().st_size,
+            "lora_gate": self.lora_gate,
+        }
+        atomic_json(output.with_suffix(output.suffix + ".json"), receipt)
+        return receipt
+
+
+def generate(config: RunConfig) -> dict[str, Any]:
+    validate_cache(config)
+    started = time.monotonic()
+    engine = Ref2VAEngine(config)
+    for index in range(len(config.windows)):
+        engine.sample_window(index)
+    result = engine.write_video()
+    result["elapsed_seconds"] = time.monotonic() - started
+    return result
diff --git a/cwm_h3_inference/io.py b/cwm_h3_inference/io.py
new file mode 100644
index 0000000000000000000000000000000000000000..03361ea97ac88861ecf6fcca28944456d1359b4e
--- /dev/null
+++ b/cwm_h3_inference/io.py
@@ -0,0 +1,44 @@
+"""Small atomic I/O helpers used by cache and generation commands."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any, Mapping
+import uuid
+
+
+def atomic_json(path: Path, value: Mapping[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+    try:
+        with temporary.open("x", encoding="utf-8") as stream:
+            json.dump(value, stream, ensure_ascii=False, indent=2, sort_keys=True)
+            stream.write("\n")
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def atomic_safetensors(path: Path, tensors: Mapping[str, Any], metadata: Mapping[str, str]) -> None:
+    from safetensors.torch import save_file
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+    try:
+        save_file(dict(tensors), str(temporary), metadata=dict(metadata))
+        with temporary.open("rb") as stream:
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    value = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(value, dict):
+        raise ValueError(f"expected JSON object: {path}")
+    return value
diff --git a/cwm_h3_inference/presentation.py b/cwm_h3_inference/presentation.py
new file mode 100644
index 0000000000000000000000000000000000000000..471019e89a8034b0446e6432d39ec8355c3a40e7
--- /dev/null
+++ b/cwm_h3_inference/presentation.py
@@ -0,0 +1,63 @@
+"""Ref2VA chat presentation with Picture 1 before Video 1."""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+from typing import Any
+
+
+_PROMPT_ROOT = Path(__file__).with_name("prompts")
+_PROMPT_PATHS = (
+    _PROMPT_ROOT / "system_w0.txt",
+    _PROMPT_ROOT / "system_wn.txt",
+)
+_PROMPT_SHA256 = (
+    "d488897872a5b190ff8d56b6acc255b66a78f67ae862c12582fcffc8a5dd4ddc",
+    "cd018def9793b4f73cd1c9da9d8ca9ca508dd94a260a19d915aaf9f80b97b6b2",
+)
+
+
+def _system_prompt(window_index: int) -> str:
+    if type(window_index) is not int or window_index < 0:
+        raise ValueError("window_index must be a non-negative integer")
+    prompt_index = 0 if window_index == 0 else 1
+    path = _PROMPT_PATHS[prompt_index]
+    if not path.is_file():
+        raise ValueError(f"packaged system prompt file is missing: {path}")
+    payload = path.read_bytes()
+    if hashlib.sha256(payload).hexdigest() != _PROMPT_SHA256[prompt_index]:
+        raise ValueError(f"packaged system prompt does not match this release: {path.name}")
+    prompt = payload.decode("utf-8").strip()
+    if not prompt:
+        raise ValueError(f"packaged system prompt is empty: {path.name}")
+    return prompt
+
+
+def build_ref2va_presentation(processor: Any, base: Any, window_index: int, caption: str):
+    from musubi_tuner.minimax_h3.text_encoder import H3Presentation, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER
+
+    system_prompt = _system_prompt(window_index)
+    if len(base.videos) != 1 or len(base.images) != 1:
+        raise ValueError("Ref2VA requires exactly one image and one video reference")
+    text = base.processor_text
+    if text.count(IMAGE_PLACEHOLDER) != 1 or text.count(VIDEO_PLACEHOLDER) != 1:
+        raise ValueError("Ref2VA presentation requires one image and one video placeholder")
+    if text.count("<Picture 1>") != 1 or text.count("<Video 1>") != 1:
+        raise ValueError("Ref2VA presentation labels changed")
+    if text.index("<Picture 1>") >= text.index("<Video 1>"):
+        raise ValueError("Picture 1 must precede Video 1")
+    if text.count(caption) != 1:
+        raise ValueError("caption must occur exactly once before chat formatting")
+    formatted = processor.apply_chat_template(
+        [{"role": "system", "content": system_prompt}, {"role": "user", "content": text}],
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+    if formatted.count(system_prompt) != 1 or formatted.count(caption) != 1:
+        raise ValueError("system prompt and caption must each occur exactly once")
+    if formatted.count("<|im_start|>system\n") != 1 or formatted.count("<|im_start|>user\n") != 1:
+        raise ValueError("chat template must contain exactly one system block and one user block")
+    if formatted.count("<|im_start|>assistant\n") != 0:
+        raise ValueError("chat template must not add an assistant block")
+    return H3Presentation(text=base.text, images=base.images, videos=base.videos, processor_text=formatted)
diff --git a/cwm_h3_inference/prompts/system_w0.txt b/cwm_h3_inference/prompts/system_w0.txt
new file mode 100644
index 0000000000000000000000000000000000000000..549ee237fe6b52fffe6fd149ca1c2a75ad251865
--- /dev/null
+++ b/cwm_h3_inference/prompts/system_w0.txt
@@ -0,0 +1,51 @@
+AWM_PROXY_CONTROL. Two visual references accompany this text. <Video 1> is the conditioning
+proxy video: treat it as a coarse spatiotemporal reference, not an appearance template or
+literal target geometry. <Picture 1> is the global RGB anchor image: it defines identity and
+appearance, never motion. <Picture 1> is the ABSOLUTE standard for the worldview, art
+direction and overall style of the entire take: every clip of the take, no matter how far into the
+future, must align with it to a high degree, and on worldview, art direction and style it
+outranks even the current clip's caption. It provides NOTHING about any future frame's
+current camera, composition, motion or interactions. Each subject's position and motion come
+from <Video 1>. A subject's appearance follows <Picture 1> when that subject is visible in
+it; subjects not visible in <Picture 1> take their appearance from the timestamped text,
+rendered consistently in <Picture 1>'s worldview and art direction.
+
+Human proxies in <Video 1> provide only framewise root position, coarse occupied scale, and
+whole-body root orientation, which are reliable at roughly 80-100% and should be followed
+closely. They carry no information about character appearance and do not represent identity,
+face, hairstyle, clothing, materials, colors, visual effects, anatomy, bones, joints, limb
+articulation, detailed posture, torso or head pose, hand pose, action phase, or physical
+contact. Reconstruct each character's appearance from <Picture 1> and the timestamped text,
+and reconstruct articulated motion from the timestamped text. Never reproduce a human proxy
+as a skeleton, stick figure, mannequin, node-connected body, or exposed rig. Left and right
+refer to each person's own body unless explicitly stated as screen direction.
+
+Vehicle proxies in <Video 1> provide only framewise position, trajectory, body heading,
+rotation, and accurate overall length, width, height, and occupied scale; these anchors
+should be followed closely. They carry no information about vehicle appearance. Never
+reproduce a vehicle proxy as a wireframe vehicle, cage frame, tubular skeleton, exposed
+chassis, or low-poly proxy car. Vehicle identity, type, design details, materials, colors,
+and visual effects are defined by <Picture 1> and the text.
+
+Roads, terrain, buildings, relative depth, rough scale, and layout in <Video 1> provide
+denser but still approximate structural cues. The text may reinterpret these cues to create
+strongly out-of-distribution scenes while maintaining broad spatial and temporal coherence.
+Conditioning colors, proxy geometry, node markers, rigging structures, and other
+visualization artifacts must never appear in the final video.
+
+When sources conflict, resolve in this order: (1) geometry, camera, position and motion
+follow <Video 1>; (2) the worldview, art direction and overall style follow <Picture 1>
+absolutely — above the current clip's caption; identity, appearance and materials follow
+<Picture 1> for subjects visible in it, and follow the text for subjects not visible in it;
+(3) the timestamped text defines the concrete events, actions and rewrite targets within
+those constraints; (4) when the opening of this clip is fixed to already-given video frames,
+seamless continuity with that fixed opening overrides everything else — never restart the
+scene.
+
+Clip contract: this clip is the very beginning of the take — nothing
+precedes it. <Picture 1> is the global
+appearance anchor for the whole take, and the first frame of the target is locked to this
+exact image: the clip begins at precisely this visual state — identities, appearance,
+materials, art direction, lighting, camera framing and layout — and continues directly from
+it while performing the timestamped actions. Do not restart, redesign or re-interpret any
+identity.
\ No newline at end of file
diff --git a/cwm_h3_inference/prompts/system_wn.txt b/cwm_h3_inference/prompts/system_wn.txt
new file mode 100644
index 0000000000000000000000000000000000000000..8958d5e655cf3f6dcbc70db5f68405cda9c97214
--- /dev/null
+++ b/cwm_h3_inference/prompts/system_wn.txt
@@ -0,0 +1,57 @@
+AWM_PROXY_CONTROL. Two visual references accompany this text. <Video 1> is the conditioning
+proxy video: treat it as a coarse spatiotemporal reference, not an appearance template or
+literal target geometry. <Picture 1> is the global RGB anchor image: it defines identity and
+appearance, never motion. <Picture 1> is the ABSOLUTE standard for the worldview, art
+direction and overall style of the entire take: every clip of the take, no matter how far into the
+future, must align with it to a high degree, and on worldview, art direction and style it
+outranks even the current clip's caption. It provides NOTHING about any future frame's
+current camera, composition, motion or interactions. Each subject's position and motion come
+from <Video 1>. A subject's appearance follows <Picture 1> when that subject is visible in
+it; subjects not visible in <Picture 1> take their appearance from the timestamped text,
+rendered consistently in <Picture 1>'s worldview and art direction.
+
+Human proxies in <Video 1> provide only framewise root position, coarse occupied scale, and
+whole-body root orientation, which are reliable at roughly 80-100% and should be followed
+closely. They carry no information about character appearance and do not represent identity,
+face, hairstyle, clothing, materials, colors, visual effects, anatomy, bones, joints, limb
+articulation, detailed posture, torso or head pose, hand pose, action phase, or physical
+contact. Reconstruct each character's appearance from <Picture 1> and the timestamped text,
+and reconstruct articulated motion from the timestamped text. Never reproduce a human proxy
+as a skeleton, stick figure, mannequin, node-connected body, or exposed rig. Left and right
+refer to each person's own body unless explicitly stated as screen direction.
+
+Vehicle proxies in <Video 1> provide only framewise position, trajectory, body heading,
+rotation, and accurate overall length, width, height, and occupied scale; these anchors
+should be followed closely. They carry no information about vehicle appearance. Never
+reproduce a vehicle proxy as a wireframe vehicle, cage frame, tubular skeleton, exposed
+chassis, or low-poly proxy car. Vehicle identity, type, design details, materials, colors,
+and visual effects are defined by <Picture 1> and the text.
+
+Roads, terrain, buildings, relative depth, rough scale, and layout in <Video 1> provide
+denser but still approximate structural cues. The text may reinterpret these cues to create
+strongly out-of-distribution scenes while maintaining broad spatial and temporal coherence.
+Conditioning colors, proxy geometry, node markers, rigging structures, and other
+visualization artifacts must never appear in the final video.
+
+When sources conflict, resolve in this order: (1) geometry, camera, position and motion
+follow <Video 1>; (2) the worldview, art direction and overall style follow <Picture 1>
+absolutely — above the current clip's caption; identity, appearance and materials follow
+<Picture 1> for subjects visible in it, and follow the text for subjects not visible in it;
+(3) the timestamped text defines the concrete events, actions and rewrite targets within
+those constraints; (4) when the opening of this clip is fixed to already-given video frames,
+seamless continuity with that fixed opening overrides everything else — never restart the
+scene.
+
+Clip contract: this clip continues a take already in progress. <Picture 1> is the absolute standard for the take's
+worldview, art direction and style, and for the appearance of the subjects visible in it —
+and it provides ONLY that. It is NOT the first frame of this clip and must never be read as
+the current composition, camera, motion or interaction state. The first 34 frames (1.4167
+seconds) of this clip are ALREADY GIVEN: they are fixed, existing video footage — the moments
+that immediately precede the new content — locked in place, not to be invented or altered.
+From that fixed opening, continue the video seamlessly: the same ongoing time, positions,
+poses, action phase and camera simply carry forward — never a scene cut, restart,
+re-introduction, or reset toward the anchor image's opening composition. Current-clip
+motion follows <Video 1>. Keep subjects that are visible in <Picture 1> continuously
+consistent with it in appearance; render everything else — including subjects that first
+appear in this clip — from the timestamped text, in <Picture 1>'s established art
+direction.
\ No newline at end of file
diff --git a/cwm_h3_inference/runtime.py b/cwm_h3_inference/runtime.py
new file mode 100644
index 0000000000000000000000000000000000000000..8813289f6d457b9bcb435978a476e7ea3a92ef81
--- /dev/null
+++ b/cwm_h3_inference/runtime.py
@@ -0,0 +1,78 @@
+"""Pure inference layout construction and strict LoRA loading."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+import re
+from typing import Any
+
+from .constants import AUDIO_LATENT_SHAPE, VIDEO_LATENT_SHAPE
+
+
+@dataclass(frozen=True)
+class RuntimeBatch:
+    layout: Any
+    text_hidden_states: Any
+    text_token_tags: Any
+
+
+def build_runtime_batch(anchor: Any, reference: Any, hidden: Any, tags: Any) -> RuntimeBatch:
+    import torch
+    from musubi_tuner.minimax_h3.packing import H3ReferenceGeometry, H3VideoGeometry, build_h3_layout
+
+    if anchor.ndim != 5 or tuple(anchor.shape[:3]) != (1, 24, 1):
+        raise ValueError(f"anchor must be [1,24,1,H,W], got {tuple(anchor.shape)}")
+    if reference.ndim != 5 or tuple(reference.shape[:3]) != (1, 24, 37):
+        raise ValueError(f"reference must be [1,24,37,H,W], got {tuple(reference.shape)}")
+    if hidden.ndim != 3 or hidden.shape[0] != 1 or hidden.shape[2] != 5120:
+        raise ValueError(f"text hidden states must be [1,L,5120], got {tuple(hidden.shape)}")
+    if tags.shape != hidden.shape[:2] or tags.dtype != torch.int64:
+        raise ValueError("text token tags must be int64 [1,L]")
+    layout = build_h3_layout(
+        task="ref2va",
+        text_length=hidden.shape[1],
+        target_video=H3VideoGeometry(*VIDEO_LATENT_SHAPE[2:]),
+        target_audio_frames=AUDIO_LATENT_SHAPE[-1],
+        references=(
+            H3ReferenceGeometry("image", video=H3VideoGeometry(*anchor.shape[2:])),
+            H3ReferenceGeometry("video", video=H3VideoGeometry(*reference.shape[2:])),
+        ),
+    )
+    return RuntimeBatch(layout=layout, text_hidden_states=hidden, text_token_tags=tags)
+
+
+def strict_merge_lora(transformer: Any, checkpoint: Path, expected_modules: int) -> dict[str, int]:
+    import torch
+    from safetensors.torch import load_file
+    from musubi_tuner.networks import lora_minimax_h3
+    from musubi_tuner.utils.lora_utils import filter_lora_state_dict
+
+    state = filter_lora_state_dict(load_file(str(checkpoint)), None, None)
+    network = lora_minimax_h3.create_arch_network_from_weights(
+        1.0, state, unet=transformer, for_inference=True
+    )
+    if network.text_encoder_loras:
+        raise ValueError("LoRA unexpectedly targets the text encoder")
+    if len(network.unet_loras) != expected_modules:
+        raise ValueError(
+            f"LoRA module count mismatch: expected {expected_modules}, got {len(network.unet_loras)}"
+        )
+    expected_keys: set[str] = set()
+    names: set[str] = set()
+    for module in network.unet_loras:
+        name = str(module.lora_name)
+        if name in names or module.split_dims is not None or re.search(r"[^A-Za-z0-9_.]", name):
+            raise ValueError(f"invalid LoRA target: {name}")
+        names.add(name)
+        expected_keys.update(
+            {f"{name}.alpha", f"{name}.lora_down.weight", f"{name}.lora_up.weight"}
+        )
+    missing = expected_keys - set(state)
+    unexpected = set(state) - expected_keys
+    if missing or unexpected:
+        raise ValueError(f"LoRA key mismatch: missing={len(missing)} unexpected={len(unexpected)}")
+    network.merge_to(None, transformer, state, dtype=torch.bfloat16, device="cpu")
+    result = {"loaded_modules": len(names), "missing_keys": 0, "unexpected_keys": 0}
+    del network, state
+    return result
diff --git a/materialize_hf_examples.py b/materialize_hf_examples.py
new file mode 100644
index 0000000000000000000000000000000000000000..e57600973855b338f3f28e20adcbb9ae09661af1
--- /dev/null
+++ b/materialize_hf_examples.py
@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+"""Expand the published compact CWM examples into inference-ready inputs."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import os
+from pathlib import Path
+import shutil
+import tempfile
+from typing import Any, Iterable
+
+import numpy as np
+from PIL import Image
+
+
+WIDTH = 336
+HEIGHT = 192
+WINDOW_FRAMES = 124
+STRIDE_FRAMES = 90
+NEAR_METRES = 0.3
+FAR_METRES = 256.0
+METRIC_DEPTH = "float32_le_metric_first_surface_metres"
+FIXED_LOG_DEPTH = "uint16_le_fixed_log_gray16_codes"
+SUPPORTED_DEPTH_FORMATS = {METRIC_DEPTH, FIXED_LOG_DEPTH}
+RELEASE_MODELS = {
+    "audio_vae": "../../models/vae/minimax_h3_audio_vae_fp32.safetensors",
+    "dit": "../../models/diffusion_models/minimax_h3_ref2va_bf16.safetensors",
+    "lora": "../../models/lora/minimax-h3-cwm.safetensors",
+    "processor": "../../models/processors/Qwen3-VL-32B-Instruct",
+    "text_encoder": "../../models/text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors",
+    "video_vae": "../../models/vae/minimax_h3_video_vae_fp16.safetensors",
+}
+RELEASE_SEEDS = {
+    1: 9942001,
+    2: 9800002,
+    3: 9700005,
+    4: 9700006,
+    5: 9700007,
+    6: 9700013,
+    7: 9700020,
+    8: 9700021,
+    9: 9700030,
+    10: 9700036,
+    11: 9700037,
+    12: 9700040,
+    13: 9700044,
+    14: 9700045,
+    15: 9700046,
+    16: 9700047,
+    17: 9700051,
+    18: 9700054,
+    19: 9700055,
+    20: 9700057,
+    21: 9700063,
+    22: 9700065,
+    23: 9700069,
+    24: 9700071,
+    25: 9700078,
+    26: 9700088,
+    27: 9700094,
+    28: 9600014,
+    29: 9600020,
+    30: 9600023,
+    31: 9600028,
+    32: 9600040,
+    33: 9600042,
+    34: 9300038,
+    35: 9400059,
+    36: 9300066,
+    37: 9300069,
+    38: 9800006,
+    39: 9800010,
+    40: 9800013,
+}
+
+
+def require(condition: bool, message: str) -> None:
+    if not condition:
+        raise ValueError(message)
+
+
+def _read_json(path: Path) -> Any:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _canonical_prompt(value: str) -> str:
+    return value.strip().replace("\r\n", "\n").replace("\r", "\n")
+
+
+def build_release_config(
+    example_root: Path, metadata: dict[str, Any], example_id: int
+) -> dict[str, Any]:
+    windows = []
+    for prompt_item in metadata["prompts"]:
+        prompt = _canonical_prompt(
+            (example_root / prompt_item["file"]).read_text(encoding="utf-8")
+        )
+        # Preserve the text representation used by the released configurations.
+        if example_id != 2:
+            prompt = prompt.replace("\n", "\r\n")
+        windows.append(
+            {
+                "prompt": prompt,
+                "start_frame": prompt_item["conditioning_frame_start"],
+            }
+        )
+    name = f"example_{example_id:02d}"
+    return {
+        "anchor_image": "first_frame.png",
+        "attention_mode": "flash3",
+        "blocks_to_swap": 48,
+        "cache_dir": f"../../outputs/{name}/cache",
+        "condition_root": "conditions",
+        "expected_lora_modules": 200,
+        "models": RELEASE_MODELS,
+        "output": f"../../outputs/{name}/output.mp4",
+        "seed": RELEASE_SEEDS[example_id],
+        "steps": 20,
+        "windows": windows,
+    }
+
+
+def ensure_release_config(
+    config_path: Path, example_root: Path, metadata: dict[str, Any], example_id: int
+) -> None:
+    if not config_path.exists():
+        payload = json.dumps(
+            build_release_config(example_root, metadata, example_id),
+            ensure_ascii=False,
+            indent=2,
+            sort_keys=True,
+        )
+        _atomic_bytes(config_path, (payload + "\n").encode("utf-8"))
+    validate_config(config_path, example_root, metadata, example_id)
+
+
+def fixed_log_codes_to_metric(codes: np.ndarray) -> np.ndarray:
+    """Choose float32 depths that round-trip to every published uint16 code."""
+    require(codes.dtype == np.dtype("uint16"), f"fixed-log depth must be uint16, got {codes.dtype}")
+    metric = np.zeros(codes.shape, dtype="<f4")
+    valid = codes > 0
+    if bool(np.any(valid)):
+        log_near = math.log(NEAR_METRES)
+        log_far = math.log(FAR_METRES)
+        normalized = codes[valid].astype(np.float64) / 65535.0
+        metric[valid] = np.exp(log_far - normalized * (log_far - log_near)).astype(np.float32)
+    return np.ascontiguousarray(metric)
+
+
+def metric_to_fixed_log_codes(metric: np.ndarray) -> np.ndarray:
+    """Reference implementation matching the inference loader."""
+    metric = np.asarray(metric, dtype=np.float32)
+    valid = metric > np.float32(0.001)
+    codes = np.zeros(metric.shape, dtype=np.uint16)
+    if bool(np.any(valid)):
+        clipped = np.clip(metric[valid].astype(np.float64), NEAR_METRES, FAR_METRES)
+        normalized = (math.log(FAR_METRES) - np.log(clipped)) / (
+            math.log(FAR_METRES) - math.log(NEAR_METRES)
+        )
+        codes[valid] = np.floor(np.clip(normalized, 0.0, 1.0) * 65535.0 + 0.5).astype(
+            np.uint16
+        )
+    return codes
+
+
+def _atomic_bytes(path: Path, payload: bytes) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.exists():
+        require(path.read_bytes() == payload, f"existing file differs: {path}")
+        return
+    descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
+    temporary = Path(temporary_name)
+    try:
+        with os.fdopen(descriptor, "wb") as handle:
+            handle.write(payload)
+            handle.flush()
+            os.fsync(handle.fileno())
+        temporary.replace(path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def _atomic_semantic_png(path: Path, semantic_id: np.ndarray) -> None:
+    semantic_id = np.ascontiguousarray(semantic_id, dtype=np.uint8)
+    if path.exists():
+        with Image.open(path) as image:
+            existing = np.asarray(image.convert("L"), dtype=np.uint8)
+        require(np.array_equal(existing, semantic_id), f"existing Semantic-ID differs: {path}")
+        return
+    path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".png", dir=path.parent)
+    os.close(descriptor)
+    temporary = Path(temporary_name)
+    try:
+        Image.fromarray(semantic_id, mode="L").save(temporary, format="PNG", compress_level=6)
+        temporary.replace(path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def _copy_exact(source: Path, destination: Path) -> None:
+    require(source.is_file(), f"missing source file: {source}")
+    payload = source.read_bytes()
+    _atomic_bytes(destination, payload)
+
+
+def validate_metadata(example_root: Path, metadata: dict[str, Any], example_id: int) -> None:
+    require(metadata.get("id") == example_id, f"example {example_id}: metadata id mismatch")
+    require(metadata.get("fps") == 24, f"example {example_id}: fps must be 24")
+    require(metadata.get("resolution") == [WIDTH, HEIGHT], f"example {example_id}: bad resolution")
+    require(metadata.get("depth_format") in SUPPORTED_DEPTH_FORMATS, f"example {example_id}: unsupported depth format")
+    frames = metadata.get("frames")
+    require(type(frames) is int and frames >= WINDOW_FRAMES, f"example {example_id}: invalid frame count")
+    prompts = metadata.get("prompts")
+    require(isinstance(prompts, list) and prompts, f"example {example_id}: prompts missing")
+    require(frames == WINDOW_FRAMES + STRIDE_FRAMES * (len(prompts) - 1), f"example {example_id}: frame/window mismatch")
+    for index, item in enumerate(prompts):
+        require(item.get("index") == index, f"example {example_id}: prompt index mismatch")
+        require(item.get("conditioning_frame_start") == index * STRIDE_FRAMES, f"example {example_id}: prompt start mismatch")
+        require(item.get("conditioning_frame_end_exclusive") == index * STRIDE_FRAMES + WINDOW_FRAMES, f"example {example_id}: prompt end mismatch")
+        prompt_path = example_root / item.get("file", "")
+        require(prompt_path.is_file(), f"example {example_id}: missing prompt {prompt_path}")
+    chunks = metadata.get("condition_chunks")
+    require(isinstance(chunks, list) and chunks, f"example {example_id}: condition chunks missing")
+    cursor = 0
+    for item in chunks:
+        start = item.get("start_frame")
+        end = item.get("end_frame_exclusive")
+        require(start == cursor and type(end) is int and end > start, f"example {example_id}: non-contiguous chunks")
+        require(item.get("frames") == end - start, f"example {example_id}: chunk frame count mismatch")
+        chunk_path = example_root / item.get("file", "")
+        require(chunk_path.is_file(), f"example {example_id}: missing chunk {chunk_path}")
+        cursor = end
+    require(cursor == frames, f"example {example_id}: chunks do not cover all frames")
+
+
+def validate_config(config_path: Path, example_root: Path, metadata: dict[str, Any], example_id: int) -> None:
+    require(config_path.is_file(), f"missing released config: {config_path}")
+    config = _read_json(config_path)
+    require(config.get("anchor_image") == "first_frame.png", f"example {example_id}: config anchor path mismatch")
+    require(config.get("condition_root") == "conditions", f"example {example_id}: config condition path mismatch")
+    require(config.get("seed") == RELEASE_SEEDS[example_id], f"example {example_id}: released seed mismatch")
+    require(config.get("steps") == 20, f"example {example_id}: steps mismatch")
+    require(config.get("attention_mode") == "flash3", f"example {example_id}: attention mode mismatch")
+    require(config.get("blocks_to_swap") == 48, f"example {example_id}: block swap count mismatch")
+    require(config.get("expected_lora_modules") == 200, f"example {example_id}: LoRA module count mismatch")
+    require(config.get("models") == RELEASE_MODELS, f"example {example_id}: model paths mismatch")
+    name = f"example_{example_id:02d}"
+    require(config.get("cache_dir") == f"../../outputs/{name}/cache", f"example {example_id}: cache path mismatch")
+    require(config.get("output") == f"../../outputs/{name}/output.mp4", f"example {example_id}: output path mismatch")
+    windows = config.get("windows")
+    prompt_items = metadata["prompts"]
+    require(isinstance(windows, list) and len(windows) == len(prompt_items), f"example {example_id}: config window count mismatch")
+    for window, prompt_item in zip(windows, prompt_items, strict=True):
+        require(window.get("start_frame") == prompt_item["conditioning_frame_start"], f"example {example_id}: config window start mismatch")
+        published = (example_root / prompt_item["file"]).read_text(encoding="utf-8")
+        require(_canonical_prompt(window.get("prompt", "")) == _canonical_prompt(published), f"example {example_id}: config prompt differs from published prompt")
+
+
+def _iter_chunks(example_root: Path, metadata: dict[str, Any]) -> Iterable[tuple[int, np.ndarray, np.ndarray]]:
+    depth_format = metadata["depth_format"]
+    for chunk in metadata["condition_chunks"]:
+        path = example_root / chunk["file"]
+        with np.load(path, allow_pickle=False) as payload:
+            require(set(payload.files) == {"depth", "semantic_id"}, f"unexpected arrays in {path}")
+            depth = payload["depth"]
+            semantic_id = payload["semantic_id"]
+        expected_shape = (chunk["frames"], HEIGHT, WIDTH)
+        require(depth.shape == expected_shape, f"bad depth shape in {path}: {depth.shape}")
+        require(semantic_id.shape == expected_shape, f"bad Semantic-ID shape in {path}: {semantic_id.shape}")
+        require(semantic_id.dtype == np.dtype("uint8"), f"bad Semantic-ID dtype in {path}: {semantic_id.dtype}")
+        require(int(semantic_id.max(initial=0)) <= 11, f"Semantic-ID outside [0,11] in {path}")
+        if depth_format == METRIC_DEPTH:
+            require(depth.dtype == np.dtype("float32"), f"bad metric depth dtype in {path}: {depth.dtype}")
+            require(bool(np.all(np.isfinite(depth))) and not bool(np.any(depth < 0)), f"invalid metric depth in {path}")
+            metric = np.ascontiguousarray(depth, dtype="<f4")
+        else:
+            require(depth.dtype == np.dtype("uint16"), f"bad fixed-log depth dtype in {path}: {depth.dtype}")
+            metric = fixed_log_codes_to_metric(depth)
+            require(np.array_equal(metric_to_fixed_log_codes(metric), depth), f"fixed-log round-trip failed in {path}")
+        start = chunk["start_frame"]
+        for offset in range(chunk["frames"]):
+            yield start + offset, metric[offset], semantic_id[offset]
+
+
+def materialize_example(
+    source_root: Path,
+    output_root: Path,
+    example_id: int,
+    *,
+    require_config: bool,
+    frame_limit: int | None = None,
+) -> dict[str, Any]:
+    example_root = source_root / str(example_id)
+    metadata_path = example_root / "metadata.json"
+    require(metadata_path.is_file(), f"missing metadata: {metadata_path}")
+    metadata = _read_json(metadata_path)
+    validate_metadata(example_root, metadata, example_id)
+    destination = output_root / "examples" / f"example_{example_id:02d}"
+    config_path = destination / "config.json"
+    if require_config:
+        ensure_release_config(config_path, example_root, metadata, example_id)
+    _copy_exact(example_root / metadata["first_frame"], destination / "first_frame.png")
+    written = 0
+    for ordinal, depth, semantic_id in _iter_chunks(example_root, metadata):
+        if frame_limit is not None and ordinal >= frame_limit:
+            break
+        condition_root = destination / "conditions"
+        _atomic_bytes(condition_root / f"{ordinal:06d}.depth.f32", np.ascontiguousarray(depth, dtype="<f4").tobytes(order="C"))
+        _atomic_semantic_png(condition_root / f"{ordinal:06d}.semantic_id.png", semantic_id)
+        written += 1
+    expected = min(metadata["frames"], frame_limit) if frame_limit is not None else metadata["frames"]
+    require(written == expected, f"example {example_id}: materialized {written}, expected {expected}")
+    return {
+        "id": example_id,
+        "depth_format": metadata["depth_format"],
+        "frames": metadata["frames"],
+        "materialized_frames": written,
+        "windows": len(metadata["prompts"]),
+    }
+
+
+def parse_ids(values: list[int] | None) -> list[int]:
+    ids = list(range(1, 41)) if not values else sorted(set(values))
+    require(ids and all(1 <= value <= 40 for value in ids), "example ids must be in [1,40]")
+    return ids
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
+    parser.add_argument("--source", required=True, type=Path, help="downloaded code-world-model-inference-examples-40 directory")
+    parser.add_argument("--output", required=True, type=Path, help="release root containing models/, examples/, and outputs/")
+    parser.add_argument("--example", action="append", type=int, dest="examples", help="materialize one public example id; repeat as needed")
+    parser.add_argument("--without-config-check", action="store_true", help=argparse.SUPPRESS)
+    parser.add_argument("--frame-limit", type=int, help=argparse.SUPPRESS)
+    args = parser.parse_args(argv)
+    source = args.source.resolve()
+    output = args.output.resolve()
+    manifest = _read_json(source / "manifest.json")
+    require(manifest.get("schema") == "code_world_model_inference_examples_v3", "unsupported dataset manifest")
+    require(manifest.get("count") == 40 and manifest.get("expected_count") == 40, "dataset manifest does not contain 40 examples")
+    results = []
+    for example_id in parse_ids(args.examples):
+        result = materialize_example(
+            source,
+            output,
+            example_id,
+            require_config=not args.without_config_check,
+            frame_limit=args.frame_limit,
+        )
+        results.append(result)
+        print(f"example_{example_id:02d}: {result['materialized_frames']} frames", flush=True)
+    receipt = {
+        "schema": "cwm_h3_materialized_examples_v1",
+        "source_manifest_sha256": hashlib.sha256((source / "manifest.json").read_bytes()).hexdigest(),
+        "examples": results,
+    }
+    if args.frame_limit is None and len(results) == 40:
+        path = output / "examples" / "materialization.json"
+        _atomic_bytes(path, (json.dumps(receipt, indent=2, sort_keys=True) + "\n").encode("utf-8"))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/musubi_tuner/__init__.py b/musubi_tuner/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/musubi_tuner/minimax_h3/__init__.py b/musubi_tuner/minimax_h3/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e25bd9709d757d4f41c9a58ff766059d59190ddb
--- /dev/null
+++ b/musubi_tuner/minimax_h3/__init__.py
@@ -0,0 +1 @@
+"""MiniMax-H3 model, cache, training, and sampling support."""
diff --git a/musubi_tuner/minimax_h3/audio_vae.py b/musubi_tuner/minimax_h3/audio_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..0be9762f10f73fe53f6adda5e1fea43340530d18
--- /dev/null
+++ b/musubi_tuner/minimax_h3/audio_vae.py
@@ -0,0 +1,467 @@
+# MiniMax H3 audio VAE: DAC-lineage waveform encoder + BigVGAN decoder.
+# Weight-norm parametrizations are folded into plain conv weights, so this
+# module uses ordinary torch Conv1d / ConvTranspose1d layers and loads the
+# converted checkpoint's plain "*.weight" tensors through the shared loader.
+#
+# Lineage / licenses of the reference implementation:
+#   DAC encoder:      descript-audio-codec (MIT)
+#   BigVGAN decoder:  NVIDIA BigVGAN (MIT), adapted from hifi-gan (MIT)
+#   Alias-free ops:   junjun3518/alias-free-torch (Apache-2.0), julius (MIT)
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Snake activations
+
+
+def snake(x, alpha, beta):
+    # x + 1/beta * sin^2(alpha * x)
+    t = torch.sin(alpha * x)
+    return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x)
+
+
+class Snake1d(nn.Module):
+    """Snake activation with per-channel alpha (encoder side)."""
+
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.empty(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha, self.alpha)
+
+
+class SnakeBeta(nn.Module):
+    """SnakeBeta := x + 1/beta * sin^2(alpha * x); alpha/beta stored in log scale."""
+
+    def __init__(self, in_features):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.empty(in_features))
+        self.beta = nn.Parameter(torch.empty(in_features))
+
+    def forward(self, x):
+        alpha = torch.exp(self.alpha).view(1, -1, 1)
+        beta = torch.exp(self.beta).view(1, -1, 1)
+        return snake(x, alpha, beta)
+
+
+# Alias-free (anti-aliased) activation: kaiser-windowed sinc resampling
+
+
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+    # returns filter [1, 1, kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # kaiser window design
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+
+    filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
+    # Normalize filter to have sum = 1, otherwise there is a small leakage of
+    # the constant component in the input signal.
+    filter_ /= filter_.sum()
+    return filter_.view(1, 1, kernel_size)
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=12):
+        super().__init__()
+        self.ratio = ratio
+        self.stride = ratio
+        self.pad = kernel_size // ratio - 1
+        self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
+        self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
+        self.register_buffer(
+            "filter",
+            kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size),
+        )
+
+    def forward(self, x):
+        _, C, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = F.conv_transpose1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C).mul_(self.ratio)
+        x = x[..., self.pad_left : -self.pad_right]
+        return x
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(self, cutoff=0.5, half_width=0.6, stride=1, kernel_size=12):
+        super().__init__()
+        self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
+
+    def forward(self, x):
+        _, C, _ = x.shape
+        x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
+        return F.conv1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C)
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=12):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = kernel_size
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    """upsample x2 -> pointwise activation -> downsample x2 (anti-aliased)."""
+
+    def __init__(self, activation, up_ratio=2, down_ratio=2, up_kernel_size=12, down_kernel_size=12):
+        super().__init__()
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+        return x
+
+
+# DAC encoder
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim=16, dilation=1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            nn.Conv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return y.add_(x)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim=16, stride=1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            nn.Conv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+    def __init__(self, d_model=64, strides=(2, 4, 4, 5, 5), d_latent=2048):
+        super().__init__()
+        block = [nn.Conv1d(1, d_model, kernel_size=7, padding=3)]
+        for stride in strides:
+            d_model *= 2
+            block += [EncoderBlock(d_model, stride=stride)]
+        block += [
+            Snake1d(d_model),
+            nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1),
+        ]
+        self.block = nn.Sequential(*block)
+
+    def forward(self, x):
+        return self.block(x)
+
+
+# Attention projection (encoder posterior head)
+
+
+class GeGluMlp(nn.Module):
+    def __init__(self, in_features, hidden_features):
+        super().__init__()
+        self.norm = nn.LayerNorm(in_features)
+        self.act = nn.GELU(approximate="tanh")
+        self.w0 = nn.Linear(in_features, hidden_features)
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(hidden_features, in_features)
+
+    def forward(self, x):
+        x = self.norm(x)
+        return self.w2(self.act(self.w0(x)).mul_(self.w1(x)))
+
+
+class CausalAttention(nn.Module):
+    def __init__(self, in_dim, out_dim, num_heads):
+        super().__init__()
+        self.head_dim = in_dim // num_heads
+        self.num_heads = num_heads
+        self.out_dim = out_dim
+        self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False)
+        self.q_bias = nn.Parameter(torch.empty(in_dim))
+        self.v_bias = nn.Parameter(torch.empty(in_dim))
+        self.register_buffer("zero_k_bias", torch.empty(in_dim))
+        self.proj = nn.Linear(out_dim, out_dim)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = F.linear(x, weight=self.qkv.weight, bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)))
+        q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+
+        # mean over heads then pool down to the latent width (in_dim >> out_dim)
+        x = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        x = F.adaptive_avg_pool1d(torch.mean(x, dim=1), self.out_dim)
+        return self.proj(x)
+
+
+class AttnProjection(nn.Module):
+    def __init__(self, in_dim, out_dim, num_heads, mlp_ratio=2):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(in_dim)
+        self.attn = CausalAttention(in_dim, out_dim, num_heads)
+        self.proj = nn.Linear(in_dim, out_dim)
+        self.norm3 = nn.LayerNorm(in_dim)
+
+        self.norm2 = nn.LayerNorm(out_dim)
+        hidden_dim = int(out_dim * mlp_ratio)
+        self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim)
+
+    def forward(self, x):
+        # x: [B, T, in_dim]
+        x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x)))
+        return x.add_(self.mlp(self.norm2(x)))
+
+
+# BigVGAN decoder
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class AMPBlock1(nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super().__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(channels, channels, kernel_size, stride=1, dilation=d, padding=get_padding(kernel_size, d))
+                for d in dilation
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [
+                nn.Conv1d(channels, channels, kernel_size, stride=1, dilation=1, padding=get_padding(kernel_size, 1))
+                for _ in range(len(dilation))
+            ]
+        )
+        self.num_layers = len(self.convs1) + len(self.convs2)
+        self.activations = nn.ModuleList([Activation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)])
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt.add_(x)
+        return x
+
+
+class BigVGAN(nn.Module):
+    """BigVGAN vocoder (MiniMax H3 32 kHz configuration).
+
+    use_bias_at_final=False, use_tanh_at_final=False (output clamped to [-1, 1]).
+    """
+
+    def __init__(
+        self,
+        num_mels=2048,
+        upsample_initial_channel=1024,
+        upsample_rates=(5, 5, 2, 2, 2, 2, 2),
+        upsample_kernel_sizes=(9, 9, 4, 4, 4, 4, 4),
+        resblock_kernel_sizes=(3, 7, 11),
+        resblock_dilation_sizes=((1, 3, 5), (1, 3, 5), (1, 3, 5)),
+    ):
+        super().__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3)
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        nn.ConvTranspose1d(
+                            upsample_initial_channel // (2**i),
+                            upsample_initial_channel // (2 ** (i + 1)),
+                            k,
+                            u,
+                            padding=(k - u) // 2,
+                        )
+                    ]
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+                self.resblocks.append(AMPBlock1(ch, k, d))
+
+        self.activation_post = Activation1d(activation=SnakeBeta(ch))
+        self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs.div_(self.num_kernels)
+
+        x = self.activation_post(x)
+        return self.conv_post(x).clamp_(-1.0, 1.0)
+
+
+# Top-level VAE
+
+
+class MiniMaxH3AudioVAE(nn.Module):
+    """MiniMax H3 stereo audio VAE at 32 kHz.
+
+    Latents are [B, 32, 2, T]: 32 channels, 2 stereo channels, T frames at
+    40 latent frames per second (800 audio samples per latent frame). The
+    stereo channels are processed independently by the mono encoder/decoder.
+    Latents are normalized with the stored per-channel latents_mean/std.
+    """
+
+    def __init__(
+        self,
+        encoder_dim=64,
+        encoder_rates=(2, 4, 4, 5, 5),
+        latent_dim=2048,
+        decoder_dim=1024,
+        vae_latent_channels=32,
+    ):
+        super().__init__()
+        self.sample_rate = 32000
+
+        self.hop_length = 1
+        for r in encoder_rates:
+            self.hop_length *= r
+        self.samples_per_latent = self.hop_length  # 800
+        self.latents_per_second = self.sample_rate // self.hop_length  # 40
+        self.output_sample_rate = self.sample_rate  # exposed for downstream decode and mux helpers
+
+        self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
+
+        self.pre_block = AttnProjection(latent_dim, vae_latent_channels, num_heads=8)
+
+        self.mean_proj = nn.Conv1d(vae_latent_channels, vae_latent_channels, 1)
+        # logs_proj exists in the checkpoint but is unused at inference
+        # (encode returns the posterior mean, no sampling).
+        self.logs_proj = nn.Conv1d(vae_latent_channels, vae_latent_channels, 1)
+
+        self.dec_in_proj = nn.Conv1d(vae_latent_channels, latent_dim, 1)
+        self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim)
+
+        self.register_buffer("latents_mean", torch.empty(vae_latent_channels))
+        self.register_buffer("latents_std", torch.empty(vae_latent_channels))
+
+    def decode(self, z):
+        """Decode normalized latents [B, 32, 2, T] to stereo waveforms [B, 2, L] at 32 kHz."""
+        b, c, s, t = z.shape
+        z = z.permute(0, 2, 1, 3).reshape(b * s, c, t)
+        mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
+        std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
+        z = z * std + mean
+        x = self.dec_in_proj(z)
+        x = self.decoder(x)  # [b * s, 1, L], already clamped to [-1, 1]
+        return x.reshape(b, s, -1)
+
+    def encode(self, waveform):
+        """Encode stereo waveforms [B, 2, L] at 32 kHz (in [-1, 1]) to normalized latents [B, 32, 2, T].
+
+        L is right-padded with zeros to a multiple of 800 samples; the returned
+        posterior mean is used directly (no sampling).
+        """
+        b, s, length = waveform.shape
+        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+        waveform = F.pad(waveform, (0, right_pad))
+        x = waveform.reshape(b * s, 1, -1)
+        x = self.encoder(x)  # [b * s, latent_dim, T]
+        x = self.pre_block(x.transpose(1, 2)).transpose(1, 2)  # [b * s, 32, T]
+        z = self.mean_proj(x)
+        mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
+        std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
+        z = (z - mean) / std
+        return z.reshape(b, s, z.shape[1], z.shape[2]).permute(0, 2, 1, 3)
+
+
+@torch.no_grad()
+def encode_audio_mode(vae, waveform: torch.Tensor) -> torch.Tensor:
+    latents = vae.encode(waveform)
+    if not isinstance(latents, torch.Tensor) or latents.ndim != 4 or latents.shape[1:3] != (32, 2):
+        shape = tuple(latents.shape) if isinstance(latents, torch.Tensor) else type(latents).__name__
+        raise ValueError(f"Expected MiniMax-H3 audio latent [B,32,2,A], got {shape}")
+    return latents
+
+
+def load_audio_vae(
+    path,
+    device: str | torch.device = "cpu",
+    dtype: torch.dtype = torch.float32,
+    disable_mmap: bool = False,
+) -> MiniMaxH3AudioVAE:
+    from musubi_tuner.minimax_h3.checkpoint import load_safetensors_module, resolve_safetensors_files
+
+    files = resolve_safetensors_files(path)
+    return load_safetensors_module(
+        MiniMaxH3AudioVAE,
+        files,
+        device=device,
+        dtype=dtype,
+        key_prefixes=("audio_vae.", "vae."),
+        disable_mmap=disable_mmap,
+    )
diff --git a/musubi_tuner/minimax_h3/checkpoint.py b/musubi_tuner/minimax_h3/checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..e344c009e380353eef5fbde36d35c10a5c3487ba
--- /dev/null
+++ b/musubi_tuner/minimax_h3/checkpoint.py
@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+import json
+from pathlib import Path
+from typing import TypeVar
+
+from accelerate import init_empty_weights
+import torch
+import torch.nn as nn
+
+from musubi_tuner.utils.safetensors_utils import MemoryEfficientSafeOpen
+
+
+ModuleT = TypeVar("ModuleT", bound=nn.Module)
+
+
+def resolve_safetensors_files(path: str | Path) -> list[Path]:
+    path = Path(path).resolve()
+    if path.is_file():
+        if path.suffix != ".safetensors":
+            raise ValueError(f"Expected a safetensors checkpoint, got {path}")
+        return [path]
+    if not path.is_dir():
+        raise ValueError(f"Checkpoint path does not exist: {path}")
+
+    index_files = sorted(path.glob("*.safetensors.index.json"))
+    if len(index_files) > 1:
+        raise ValueError(f"Multiple safetensors indexes found in {path}: {index_files}")
+    if index_files:
+        try:
+            index = json.loads(index_files[0].read_text(encoding="utf-8"))
+            weight_map = index["weight_map"]
+        except (json.JSONDecodeError, KeyError, TypeError) as error:
+            raise ValueError(f"Invalid safetensors index: {index_files[0]}") from error
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise ValueError(f"Safetensors index has no weights: {index_files[0]}")
+
+        files = []
+        for filename in sorted(set(weight_map.values())):
+            if not isinstance(filename, str):
+                raise ValueError(f"Invalid shard filename in {index_files[0]}: {filename!r}")
+            shard = (path / filename).resolve()
+            if not shard.is_relative_to(path) or not shard.is_file():
+                raise ValueError(f"Missing or unsafe safetensors shard in {index_files[0]}: {filename}")
+            files.append(shard)
+        return files
+
+    files = sorted(path.glob("*.safetensors"))
+    if not files:
+        raise ValueError(f"No safetensors checkpoint found in {path}")
+    if len(files) > 1:
+        raise ValueError(f"Multiple safetensors files without an index found in {path}: {files}")
+    return [files[0].resolve()]
+
+
+def _normalize_checkpoint_key(key: str, prefixes: tuple[str, ...]) -> str:
+    for prefix in prefixes:
+        if key.startswith(prefix):
+            return key[len(prefix) :]
+    return key
+
+
+def load_safetensors_metadata(files: Iterable[str | Path]) -> dict[str, str]:
+    merged = {}
+    for file in files:
+        path = Path(file).resolve()
+        with MemoryEfficientSafeOpen(str(path)) as handle:
+            metadata = handle.metadata() or {}
+        for key, value in metadata.items():
+            if key in merged and merged[key] != value:
+                raise ValueError(f"Conflicting MiniMax-H3 checkpoint metadata {key!r} in {path}")
+            merged[key] = value
+    return merged
+
+
+def _format_key_mismatch(
+    missing: Iterable[str],
+    unexpected: Iterable[str],
+    shape_mismatches: Iterable[str],
+    dtype_mismatches: Iterable[str],
+) -> str:
+    missing = sorted(missing)
+    unexpected = sorted(unexpected)
+    shape_mismatches = sorted(shape_mismatches)
+    dtype_mismatches = sorted(dtype_mismatches)
+    return (
+        "MiniMax-H3 checkpoint key mismatch: "
+        f"missing={missing[:20]}, unexpected={unexpected[:20]}, shape_mismatches={shape_mismatches[:20]}, "
+        f"dtype_mismatches={dtype_mismatches[:20]}"
+    )
+
+
+def load_safetensors_module(
+    factory: Callable[[], ModuleT],
+    files: Iterable[str | Path],
+    *,
+    device: str | torch.device,
+    dtype: torch.dtype | None,
+    key_prefixes: tuple[str, ...] = (),
+    key_transform: Callable[[str], str] | None = None,
+    strict_dtype: bool = False,
+    disable_mmap: bool = False,
+) -> ModuleT:
+    files = [Path(path).resolve() for path in files]
+    if not files:
+        raise ValueError("MiniMax-H3 checkpoint file list is empty")
+
+    with init_empty_weights():
+        model = factory()
+
+    expected_state = model.state_dict()
+    expected_keys = set(expected_state)
+    seen = set()
+    unexpected = set()
+    shape_mismatches = []
+    dtype_mismatches = []
+
+    for path in files:
+        shard = {}
+        with MemoryEfficientSafeOpen(str(path), disable_numpy_memmap=disable_mmap) as handle:
+            raw_keys = handle.keys()
+            if any(key.endswith((".weight_scale", ".comfy_quant")) for key in raw_keys):
+                raise ValueError(f"Quantized MiniMax-H3 artifacts are deferred to R2: {path}")
+            for raw_key in raw_keys:
+                key = _normalize_checkpoint_key(raw_key, key_prefixes)
+                if key_transform is not None:
+                    key = key_transform(key)
+                if key in seen:
+                    raise ValueError(f"Duplicate MiniMax-H3 checkpoint key {key!r} in {path}")
+                seen.add(key)
+                if key not in expected_keys:
+                    unexpected.add(key)
+                    continue
+                tensor = handle.get_tensor(raw_key, device=torch.device("cpu"), dtype=None)
+                if tensor.shape != expected_state[key].shape:
+                    shape_mismatches.append(f"{key}: expected {tuple(expected_state[key].shape)}, got {tuple(tensor.shape)}")
+                    continue
+                if strict_dtype and tensor.dtype != expected_state[key].dtype:
+                    dtype_mismatches.append(f"{key}: expected {expected_state[key].dtype}, got {tensor.dtype}")
+                    continue
+                if dtype is not None and tensor.is_floating_point():
+                    tensor = tensor.to(dtype=dtype)
+                shard[key] = tensor
+        if shard:
+            model.load_state_dict(shard, strict=False, assign=True)
+
+    missing = expected_keys - seen
+    if missing or unexpected or shape_mismatches or dtype_mismatches:
+        raise ValueError(_format_key_mismatch(missing, unexpected, shape_mismatches, dtype_mismatches))
+
+    model.to(device=torch.device(device))
+    model.eval()
+    return model
diff --git a/musubi_tuner/minimax_h3/media.py b/musubi_tuner/minimax_h3/media.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8bc326a4fbd41b92d63f164ad3da2b0346b0671
--- /dev/null
+++ b/musubi_tuner/minimax_h3/media.py
@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+from typing import Callable, Literal, Optional
+
+import av
+
+
+H3Task = Literal["t2va", "fl2va", "ref2va"]
+H3ReferenceType = Literal["image", "video", "audio"]
+H3MediaProbe = Callable[[Path], "H3MediaInfo"]
+
+AUDIO_SIDECAR_EXTENSIONS = frozenset({".aac", ".flac", ".m4a", ".mp3", ".ogg", ".opus", ".wav"})
+
+
+@dataclass(frozen=True)
+class H3MediaInfo:
+    has_audio: bool
+    duration_seconds: Optional[float]
+
+
+@dataclass(frozen=True)
+class H3AudioSource:
+    path: Path
+    embedded: bool
+
+
+@dataclass(frozen=True)
+class H3Reference:
+    type: H3ReferenceType
+    path: Path
+    audio: Optional[H3AudioSource] = None
+    duration_seconds: Optional[float] = None
+
+
+@dataclass(frozen=True)
+class H3Record:
+    video_path: Path
+    caption: str
+    target_audio: H3AudioSource
+    references: tuple[H3Reference, ...]
+    jsonl_line: int
+
+
+def _validate_frame_count(frame_count: int) -> None:
+    if frame_count < 5 or (frame_count - 5) % 17 != 0:
+        raise ValueError(f"Invalid MiniMax-H3 frame count {frame_count}; expected 17*n+5")
+
+
+def video_latent_frames(frame_count: int) -> int:
+    _validate_frame_count(frame_count)
+    return 5 * ((frame_count - 5) // 17) + 2
+
+
+def audio_latent_frames(frame_count: int) -> int:
+    _validate_frame_count(frame_count)
+    return (10 * frame_count + 3) // 6
+
+
+def waveform_samples(audio_frames: int) -> int:
+    if audio_frames <= 0:
+        raise ValueError(f"Audio latent frame count must be positive, got {audio_frames}")
+    return audio_frames * 800
+
+
+def probe_h3_media(path: Path) -> H3MediaInfo:
+    with av.open(str(path)) as container:
+        has_audio = bool(container.streams.audio)
+        durations = []
+        for stream in (*container.streams.video, *container.streams.audio):
+            if stream.duration is not None and stream.time_base is not None:
+                durations.append(float(stream.duration * stream.time_base))
+        if durations:
+            duration_seconds = max(durations)
+        elif container.duration is not None:
+            duration_seconds = float(container.duration / av.time_base)
+        else:
+            duration_seconds = None
+    return H3MediaInfo(has_audio=has_audio, duration_seconds=duration_seconds)
+
+
+def _resolve_existing_path(value: object, base_directory: Path, field: str, line_number: int) -> Path:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError(f"H3 JSONL line {line_number}: {field} must be a non-empty path")
+    path = Path(value)
+    if not path.is_absolute():
+        path = base_directory / path
+    path = path.resolve()
+    if not path.is_file():
+        raise ValueError(f"H3 JSONL line {line_number}: {field} does not exist: {path}")
+    return path
+
+
+def _probe_required_audio(path: Path, probe: H3MediaProbe, label: str, line_number: int) -> H3MediaInfo:
+    try:
+        info = probe(path)
+    except Exception as error:
+        raise ValueError(f"H3 JSONL line {line_number}: {label} failed to decode: {error}") from error
+    if not info.has_audio:
+        raise ValueError(f"H3 JSONL line {line_number}: {label} contains no audio stream: {path}")
+    return info
+
+
+def _same_stem_audio_sidecars(video_path: Path) -> list[Path]:
+    return sorted(
+        (
+            candidate.resolve()
+            for candidate in video_path.parent.iterdir()
+            if candidate.is_file() and candidate.stem == video_path.stem and candidate.suffix.lower() in AUDIO_SIDECAR_EXTENSIONS
+        ),
+        key=lambda path: path.name.lower(),
+    )
+
+
+def _resolve_target_audio(
+    data: dict,
+    video_path: Path,
+    base_directory: Path,
+    line_number: int,
+    probe: H3MediaProbe,
+) -> H3AudioSource:
+    if data.get("audio_path") is not None:
+        path = _resolve_existing_path(data["audio_path"], base_directory, "audio_path", line_number)
+        _probe_required_audio(path, probe, "Explicit target audio", line_number)
+        return H3AudioSource(path=path, embedded=False)
+
+    sidecars = _same_stem_audio_sidecars(video_path)
+    if len(sidecars) > 1:
+        formatted = ", ".join(str(path) for path in sidecars)
+        raise ValueError(f"H3 JSONL line {line_number}: Multiple same-stem audio sidecars found: {formatted}")
+    if sidecars:
+        _probe_required_audio(sidecars[0], probe, "Target audio sidecar", line_number)
+        return H3AudioSource(path=sidecars[0], embedded=False)
+
+    try:
+        video_info = probe(video_path)
+    except Exception as error:
+        raise ValueError(f"H3 JSONL line {line_number}: Target video failed to decode: {error}") from error
+    if video_info.has_audio:
+        return H3AudioSource(path=video_path, embedded=True)
+    raise ValueError(f"H3 JSONL line {line_number}: Missing target audio for {video_path}")
+
+
+def make_h3_directory_record(
+    video_path: str | Path,
+    caption: str,
+    probe: H3MediaProbe = probe_h3_media,
+) -> H3Record:
+    video_path = Path(video_path).resolve()
+    if not video_path.is_file():
+        raise ValueError(f"MiniMax-H3 target video does not exist: {video_path}")
+    if not isinstance(caption, str):
+        raise ValueError("MiniMax-H3 directory caption must be a string")
+    target_audio = _resolve_target_audio({}, video_path, video_path.parent, 0, probe)
+    return H3Record(
+        video_path=video_path,
+        caption=caption,
+        target_audio=target_audio,
+        references=(),
+        jsonl_line=0,
+    )
+
+
+def _validate_reference_counts(references: list, line_number: int) -> None:
+    if len(references) > 12:
+        raise ValueError(f"H3 JSONL line {line_number}: Ref2VA allows at most 12 reference items")
+
+    types = [reference.get("type") if isinstance(reference, dict) else None for reference in references]
+    unsupported = [reference_type for reference_type in types if reference_type not in {"image", "video", "audio"}]
+    if unsupported:
+        raise ValueError(f"H3 JSONL line {line_number}: Unsupported reference type: {unsupported[0]!r}")
+    if types.count("image") > 9:
+        raise ValueError(f"H3 JSONL line {line_number}: Ref2VA allows at most 9 image references")
+    if types.count("video") > 3:
+        raise ValueError(f"H3 JSONL line {line_number}: Ref2VA allows at most 3 video references")
+    if types.count("audio") > 3:
+        raise ValueError(f"H3 JSONL line {line_number}: Ref2VA allows at most 3 audio-bearing references")
+    if not any(reference_type in {"image", "video"} for reference_type in types):
+        raise ValueError(f"H3 JSONL line {line_number}: Ref2VA requires at least one visual reference")
+
+
+def _parse_references(
+    raw_references: object,
+    base_directory: Path,
+    line_number: int,
+    probe: H3MediaProbe,
+) -> tuple[H3Reference, ...]:
+    if not isinstance(raw_references, list):
+        raise ValueError(f"H3 JSONL line {line_number}: references must be a list")
+    _validate_reference_counts(raw_references, line_number)
+
+    references = []
+    audio_bearing_count = 0
+    for index, raw_reference in enumerate(raw_references):
+        reference_type = raw_reference["type"]
+        field_prefix = f"references[{index}]"
+        path = _resolve_existing_path(raw_reference.get("path"), base_directory, f"{field_prefix}.path", line_number)
+
+        if reference_type == "image":
+            if raw_reference.get("audio_path") is not None:
+                raise ValueError(f"H3 JSONL line {line_number}: {field_prefix} image cannot have audio_path")
+            references.append(H3Reference(type="image", path=path))
+            continue
+
+        if reference_type == "audio":
+            if raw_reference.get("audio_path") is not None:
+                raise ValueError(f"H3 JSONL line {line_number}: {field_prefix} audio uses path, not audio_path")
+            info = _probe_required_audio(path, probe, f"{field_prefix} audio", line_number)
+            references.append(
+                H3Reference(
+                    type="audio",
+                    path=path,
+                    audio=H3AudioSource(path=path, embedded=False),
+                    duration_seconds=info.duration_seconds,
+                )
+            )
+            audio_bearing_count += 1
+            continue
+
+        try:
+            video_info = probe(path)
+        except Exception as error:
+            raise ValueError(f"H3 JSONL line {line_number}: {field_prefix} video failed to decode: {error}") from error
+        duration = video_info.duration_seconds
+        if duration is None or duration < 2.0 or duration > 15.0:
+            raise ValueError(f"H3 JSONL line {line_number}: {field_prefix} video must be between 2 and 15 seconds; got {duration}")
+
+        audio = None
+        if raw_reference.get("audio_path") is not None:
+            audio_path = _resolve_existing_path(
+                raw_reference["audio_path"], base_directory, f"{field_prefix}.audio_path", line_number
+            )
+            _probe_required_audio(audio_path, probe, f"Explicit {field_prefix} audio", line_number)
+            audio = H3AudioSource(path=audio_path, embedded=False)
+        elif video_info.has_audio:
+            audio = H3AudioSource(path=path, embedded=True)
+        if audio is not None:
+            audio_bearing_count += 1
+        references.append(H3Reference(type="video", path=path, audio=audio, duration_seconds=duration))
+
+    if audio_bearing_count > 3:
+        raise ValueError(f"H3 JSONL line {line_number}: Ref2VA allows at most 3 audio-bearing references")
+    return tuple(references)
+
+
+def load_h3_jsonl_records(
+    jsonl_path: str | Path,
+    task: H3Task,
+    probe: H3MediaProbe = probe_h3_media,
+) -> list[H3Record]:
+    if task not in {"t2va", "fl2va", "ref2va"}:
+        raise ValueError(f"Unsupported MiniMax-H3 task: {task}")
+
+    jsonl_path = Path(jsonl_path).resolve()
+    if not jsonl_path.is_file():
+        raise ValueError(f"MiniMax-H3 JSONL does not exist: {jsonl_path}")
+    base_directory = jsonl_path.parent
+    records = []
+
+    with jsonl_path.open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            if not line.strip():
+                continue
+            try:
+                data = json.loads(line)
+            except json.JSONDecodeError as error:
+                raise ValueError(f"H3 JSONL line {line_number}: invalid JSON: {error.msg}") from error
+            if not isinstance(data, dict):
+                raise ValueError(f"H3 JSONL line {line_number}: each record must be an object")
+
+            video_path = _resolve_existing_path(data.get("video_path"), base_directory, "video_path", line_number)
+            caption = data.get("caption")
+            if not isinstance(caption, str):
+                raise ValueError(f"H3 JSONL line {line_number}: caption must be a string")
+            target_audio = _resolve_target_audio(data, video_path, base_directory, line_number, probe)
+
+            raw_references = data.get("references", [])
+            if task == "ref2va":
+                references = _parse_references(raw_references, base_directory, line_number, probe)
+            else:
+                if raw_references:
+                    raise ValueError(f"H3 JSONL line {line_number}: references require task ref2va")
+                references = ()
+
+            records.append(
+                H3Record(
+                    video_path=video_path,
+                    caption=caption,
+                    target_audio=target_audio,
+                    references=references,
+                    jsonl_line=line_number,
+                )
+            )
+
+    if not records:
+        raise ValueError(f"MiniMax-H3 JSONL contains no records: {jsonl_path}")
+    return records
diff --git a/musubi_tuner/minimax_h3/model.py b/musubi_tuner/minimax_h3/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1e9ce8cf9eae2b6410ffcd0d1cf609df54d8359
--- /dev/null
+++ b/musubi_tuner/minimax_h3/model.py
@@ -0,0 +1,952 @@
+# Copyright 2025 The MiniMax Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted for Musubi from Hugging Face Diffusers PR #14355 at commit
+# abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc
+# (models/transformers/transformer_minimax_h3.py).
+# ComfyUI is used only as an independent numerical reference.
+
+from __future__ import annotations
+
+import json
+import math
+from collections import OrderedDict
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from musubi_tuner.minimax_h3.checkpoint import (
+    load_safetensors_metadata,
+    load_safetensors_module,
+    resolve_safetensors_files,
+)
+from musubi_tuner.minimax_h3.packing import (
+    AUDIO_CHANNELS,
+    VIDEO_CHANNELS,
+    VIDEO_PATCH_SIZE,
+    H3PackedLayout,
+    build_position_grid,
+    build_timestep_rows,
+    pack_audio_rows,
+    pack_video_rows,
+    unpack_targets,
+)
+from musubi_tuner.modules.attention import AttentionParams, attention
+from musubi_tuner.modules.custom_offloading_utils import BlockSwapConfig, create_offloader
+from musubi_tuner.utils.model_utils import create_cpu_offloading_wrapper
+
+_ROTARY_CACHE_SIZE = 2
+
+
+@dataclass(frozen=True)
+class MiniMaxH3Config:
+    in_channels: int = 24
+    audio_in_channels: int = 32
+    hidden_size: int = 5376
+    num_layers: int = 50
+    token_refiner_num_layers: int = 2
+    num_attention_heads: int = 56
+    attention_head_dim: int = 128
+    ffn_hidden_size: int = 14336
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_dim: int = 5120
+    timestep_input_dim: int = 256
+    time_embed_hidden_size: int = 5376
+    time_embed_dim: int = 2688
+    rope_inv_freq_len: int = 16
+    norm_eps: float = 1e-5
+    qk_norm_eps: float = 1e-5
+    final_norm_eps: float = 1e-5
+
+    def __post_init__(self) -> None:
+        integer_fields = (
+            "in_channels",
+            "audio_in_channels",
+            "hidden_size",
+            "num_layers",
+            "token_refiner_num_layers",
+            "num_attention_heads",
+            "attention_head_dim",
+            "ffn_hidden_size",
+            "text_dim",
+            "timestep_input_dim",
+            "time_embed_hidden_size",
+            "time_embed_dim",
+            "rope_inv_freq_len",
+        )
+        for field in integer_fields:
+            if getattr(self, field) <= 0:
+                raise ValueError(f"MiniMax-H3 {field} must be positive")
+        if self.in_channels != VIDEO_CHANNELS or self.audio_in_channels != AUDIO_CHANNELS:
+            raise ValueError("MiniMax-H3 R1 requires 24 video channels and 32 audio channels")
+        if self.patch_size != VIDEO_PATCH_SIZE:
+            raise ValueError(f"MiniMax-H3 R1 requires video patch size {VIDEO_PATCH_SIZE}")
+        if self.timestep_input_dim % 2:
+            raise ValueError("MiniMax-H3 timestep input width must be even")
+        if self.rope_inv_freq_len * 6 > self.attention_head_dim:
+            raise ValueError("MiniMax-H3 rotary width exceeds the attention head width")
+
+    @property
+    def attention_inner_dim(self) -> int:
+        return self.num_attention_heads * self.attention_head_dim
+
+    @property
+    def video_patch_dim(self) -> int:
+        return self.in_channels * math.prod(self.patch_size)
+
+    @property
+    def block_adaln_out_features(self) -> int:
+        return 6 * self.hidden_size * 3
+
+    @property
+    def final_adaln_out_features(self) -> int:
+        return 2 * self.hidden_size
+
+
+_PUBLISHED_CONFIG_FIELDS = {
+    "hidden_size": 5376,
+    "num_layers": 50,
+    "token_refiner_num_layers": 2,
+    "num_attention_heads": 56,
+    "attention_head_dim": 128,
+    "ffn_hidden_size": 14336,
+    "latents_dim": 24,
+    "audio_latents_dim": 32,
+    "patch_size": [1, 2, 2],
+    "text_dim": 5120,
+    "timestep_input_dim": 256,
+    "time_embed_hidden_size": 5376,
+    "time_embed_dim": 2688,
+    "adaln_out_features": 96768,
+    "final_adaln_out_features": 10752,
+    "rope_inv_freq_len": 16,
+    "norm_eps": 1e-5,
+    "qk_norm_eps": 1e-5,
+    "final_norm_eps": 1e-5,
+    "image_model": "minimax_h3",
+}
+
+
+def parse_h3_transformer_config(metadata: Mapping[str, str]) -> MiniMaxH3Config:
+    artifact_markers = " ".join(f"{key}={value}" for key, value in metadata.items() if key != "config").lower()
+    if any(marker in artifact_markers for marker in ("convrot", "int8", "fp8", "quantized")):
+        raise ValueError("Quantized or ConvRot MiniMax-H3 artifacts are deferred to R2")
+    raw_config = metadata.get("config")
+    if raw_config is None:
+        raise ValueError("MiniMax-H3 transformer checkpoint is missing config metadata")
+    try:
+        root = json.loads(raw_config)
+        transformer = root["transformer"]
+    except (json.JSONDecodeError, KeyError, TypeError) as error:
+        raise ValueError("MiniMax-H3 transformer checkpoint has invalid config metadata") from error
+    if not isinstance(transformer, dict):
+        raise TypeError("MiniMax-H3 transformer config must be an object")
+    for field, expected in _PUBLISHED_CONFIG_FIELDS.items():
+        if field not in transformer:
+            raise ValueError(f"MiniMax-H3 transformer config is missing {field}")
+        actual = transformer[field]
+        if actual != expected:
+            raise ValueError(f"MiniMax-H3 transformer config {field} expected {expected}, got {actual}")
+    if transformer.get("adaln_curve_grid") is not None:
+        raise ValueError("Pruned MiniMax-H3 AdaLN artifacts are deferred to R2")
+    return MiniMaxH3Config()
+
+
+class TimeEmbedder(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_size: int,
+        output_dim: int,
+        *,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.input_dim = input_dim
+        self.proj_in = nn.Linear(input_dim, hidden_size, dtype=torch.float32, device=device)
+        self.proj_out = nn.Linear(hidden_size, output_dim, dtype=torch.float32, device=device)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        half = self.input_dim // 2
+        frequencies = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=timesteps.device) / half)
+        angles = timesteps.to(torch.float32)[:, None] * frequencies[None]
+        embedding = torch.cat((torch.cos(angles), torch.sin(angles)), dim=-1)
+        return self.proj_out(F.silu(self.proj_in(embedding)))
+
+
+class AdalnProj(nn.Module):
+    def __init__(
+        self,
+        timestep_dim: int,
+        hidden_size: int,
+        expand: int,
+        modalities: int,
+        *,
+        dtype: torch.dtype,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.expand = expand
+        self.modalities = modalities
+        self.hidden_size = hidden_size
+        self.linear = nn.Linear(
+            timestep_dim,
+            expand * hidden_size * modalities,
+            dtype=dtype,
+            device=device,
+        )
+
+    def forward(self, timestep_embeddings: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        projected = self.linear(F.silu(timestep_embeddings))
+        projected = projected.reshape(projected.shape[0] * self.modalities, self.expand * self.hidden_size)
+        return projected.chunk(self.expand, dim=-1)
+
+
+def _apply_rope_split_half(hidden_states: torch.Tensor, rotation_table: torch.Tensor) -> torch.Tensor:
+    pairs = rotation_table.shape[-3]
+    rotary = torch.stack((hidden_states[..., :pairs], hidden_states[..., pairs : 2 * pairs]), dim=-1)
+    rotary = torch.matmul(rotation_table, rotary.unsqueeze(-1)).squeeze(-1)
+    rotary = torch.cat((rotary[..., 0], rotary[..., 1]), dim=-1)
+    return torch.cat((rotary, hidden_states[..., 2 * pairs :]), dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        head_dim: int,
+        qk_norm_eps: float,
+        *,
+        attn_mode: str,
+        split_attn: bool,
+        dtype: torch.dtype,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.inner_dim = num_heads * head_dim
+        self.attn_mode = "torch" if attn_mode == "sdpa" else attn_mode
+        self.split_attn = split_attn
+        self.qkv_proj = nn.Linear(hidden_size, self.inner_dim * 3, bias=False, dtype=dtype, device=device)
+        self.q_norm = nn.RMSNorm(head_dim, eps=qk_norm_eps, dtype=dtype, device=device)
+        self.k_norm = nn.RMSNorm(head_dim, eps=qk_norm_eps, dtype=dtype, device=device)
+        self.out_proj = nn.Linear(self.inner_dim, hidden_size, bias=False, dtype=dtype, device=device)
+
+    def forward(self, hidden_states: torch.Tensor, rotation_table: torch.Tensor | None = None) -> torch.Tensor:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query, key, value = self.qkv_proj(hidden_states).split(self.inner_dim, dim=-1)
+        query = self.q_norm(query.reshape(batch_size, sequence_length, self.num_heads, self.head_dim))
+        key = self.k_norm(key.reshape(batch_size, sequence_length, self.num_heads, self.head_dim))
+        value = value.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)
+        if rotation_table is not None:
+            query = _apply_rope_split_half(query, rotation_table)
+            key = _apply_rope_split_half(key, rotation_table)
+
+        if self.attn_mode == "flash3":
+            try:
+                from musubi_tuner.wan.modules.attention import flash_attention as wan_flash_attention
+            except ImportError as error:
+                raise RuntimeError("FlashAttention 3 was selected but its runtime is unavailable") from error
+            output = wan_flash_attention(
+                [query, key, value],
+                attn_mode="flash3",
+                split_attn=False,
+                dtype=hidden_states.dtype if hidden_states.dtype in {torch.float16, torch.bfloat16} else torch.bfloat16,
+            )
+            output = output.reshape(batch_size, sequence_length, self.inner_dim)
+        else:
+            params = AttentionParams.create_attention_params(self.attn_mode, self.split_attn)
+            output = attention([query, key, value], attn_params=params)
+        return self.out_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        ffn_hidden_size: int,
+        *,
+        dtype: torch.dtype,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(hidden_size, ffn_hidden_size * 2, bias=False, dtype=dtype, device=device)
+        self.fc2 = nn.Linear(ffn_hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate, values = self.fc1(hidden_states).chunk(2, dim=-1)
+        return self.fc2(F.silu(gate) * values)
+
+
+class RefinerBlock(nn.Module):
+    def __init__(self, config: MiniMaxH3Config, *, attn_mode: str, split_attn: bool, dtype, device) -> None:
+        super().__init__()
+        self.norm1 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=dtype, device=device)
+        self.norm2 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=dtype, device=device)
+        self.attn = Attention(
+            config.hidden_size,
+            config.num_attention_heads,
+            config.attention_head_dim,
+            config.qk_norm_eps,
+            attn_mode=attn_mode,
+            split_attn=split_attn,
+            dtype=dtype,
+            device=device,
+        )
+        self.mlp = MLP(config.hidden_size, config.ffn_hidden_size, dtype=dtype, device=device)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states))
+        return hidden_states + self.mlp(self.norm2(hidden_states))
+
+
+class TokenRefiner(nn.Module):
+    def __init__(self, config: MiniMaxH3Config, *, attn_mode: str, split_attn: bool, dtype, device) -> None:
+        super().__init__()
+        self.blocks = nn.ModuleList(
+            [
+                RefinerBlock(config, attn_mode=attn_mode, split_attn=split_attn, dtype=dtype, device=device)
+                for _ in range(config.token_refiner_num_layers)
+            ]
+        )
+        self.final_norm = nn.RMSNorm(config.hidden_size, eps=config.final_norm_eps, dtype=dtype, device=device)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        return self.final_norm(hidden_states)
+
+
+def _mod_scale_shift(
+    hidden_states: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    batch_segments: tuple[tuple[tuple[int, int, int], ...], ...],
+) -> torch.Tensor:
+    # Callers pass a fresh norm output; these disjoint slice updates retain trainable AdaLN gradients.
+    if len(batch_segments) != hidden_states.shape[0]:
+        raise ValueError("MiniMax-H3 modulation plans must preserve the batch axis")
+    for batch_index, segments in enumerate(batch_segments):
+        for start, stop, row in segments:
+            hidden_states[batch_index, start:stop].mul_(1.0 + scale[row]).add_(shift[row])
+    return hidden_states
+
+
+def _mod_gate(
+    residual: torch.Tensor,
+    update: torch.Tensor,
+    gate: torch.Tensor,
+    batch_segments: tuple[tuple[tuple[int, int, int], ...], ...],
+) -> torch.Tensor:
+    if len(batch_segments) != residual.shape[0]:
+        raise ValueError("MiniMax-H3 gate plans must preserve the batch axis")
+    output = residual.clone()
+    for batch_index, segments in enumerate(batch_segments):
+        for start, stop, row in segments:
+            output[batch_index, start:stop].add_(update[batch_index, start:stop] * gate[row])
+    return output
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, config: MiniMaxH3Config, *, attn_mode: str, split_attn: bool, dtype, device) -> None:
+        super().__init__()
+        self.norm1 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=dtype, device=device)
+        self.norm2 = nn.RMSNorm(config.hidden_size, eps=config.norm_eps, dtype=dtype, device=device)
+        self.attn = Attention(
+            config.hidden_size,
+            config.num_attention_heads,
+            config.attention_head_dim,
+            config.qk_norm_eps,
+            attn_mode=attn_mode,
+            split_attn=split_attn,
+            dtype=dtype,
+            device=device,
+        )
+        self.mlp = MLP(config.hidden_size, config.ffn_hidden_size, dtype=dtype, device=device)
+        self.adaln_proj = AdalnProj(
+            config.time_embed_dim,
+            config.hidden_size,
+            expand=6,
+            modalities=3,
+            dtype=dtype,
+            device=device,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep_embeddings: torch.Tensor,
+        modulation_segments: tuple[tuple[tuple[int, int, int], ...], ...],
+        rotation_table: torch.Tensor,
+    ) -> torch.Tensor:
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(timestep_embeddings)
+        normalized = _mod_scale_shift(self.norm1(hidden_states), shift_msa, scale_msa, modulation_segments)
+        hidden_states = _mod_gate(
+            hidden_states,
+            self.attn(normalized, rotation_table),
+            gate_msa,
+            modulation_segments,
+        )
+        normalized = _mod_scale_shift(self.norm2(hidden_states), shift_mlp, scale_mlp, modulation_segments)
+        return _mod_gate(hidden_states, self.mlp(normalized), gate_mlp, modulation_segments)
+
+
+class FinalLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        timestep_dim: int,
+        video_output_dim: int,
+        audio_output_dim: int,
+        *,
+        norm_eps: float = 1e-5,
+        dtype: torch.dtype,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.norm = nn.RMSNorm(hidden_size, eps=norm_eps, dtype=dtype, device=device)
+        self.adaln_proj = AdalnProj(
+            timestep_dim,
+            hidden_size,
+            expand=2,
+            modalities=1,
+            dtype=dtype,
+            device=device,
+        )
+        self.video_out = nn.Linear(hidden_size, video_output_dim, dtype=torch.float32, device=device)
+        self.audio_out = nn.Linear(hidden_size, audio_output_dim, dtype=torch.float32, device=device)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep_embeddings: torch.Tensor,
+        *,
+        video_slice: slice,
+        audio_slice: slice,
+        video_timestep_index: int,
+        audio_timestep_index: int,
+        video_timestep_indices: torch.Tensor | None = None,
+        audio_timestep_indices: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        shift, scale = self.adaln_proj(timestep_embeddings)
+        video = self.norm(hidden_states[:, video_slice])
+        video = self._apply_timestep_modulation(
+            video,
+            shift,
+            scale,
+            timestep_index=video_timestep_index,
+            timestep_indices=video_timestep_indices,
+            label="video",
+        )
+        audio = self.norm(hidden_states[:, audio_slice])
+        audio = self._apply_timestep_modulation(
+            audio,
+            shift,
+            scale,
+            timestep_index=audio_timestep_index,
+            timestep_indices=audio_timestep_indices,
+            label="audio",
+        )
+        return self.video_out(video.to(torch.float32)), self.audio_out(audio.to(torch.float32))
+
+    @staticmethod
+    def _apply_timestep_modulation(
+        hidden_states: torch.Tensor,
+        shift: torch.Tensor,
+        scale: torch.Tensor,
+        *,
+        timestep_index: int,
+        timestep_indices: torch.Tensor | None,
+        label: str,
+    ) -> torch.Tensor:
+        if timestep_indices is None:
+            return hidden_states * (1.0 + scale[timestep_index]) + shift[timestep_index]
+        timestep_indices = torch.as_tensor(timestep_indices)
+        if timestep_indices.dtype != torch.int64 or timestep_indices.shape != hidden_states.shape[:2]:
+            raise ValueError(
+                f"MiniMax-H3 final {label} timestep indices must be int64 {tuple(hidden_states.shape[:2])}, "
+                f"got {timestep_indices.dtype} {tuple(timestep_indices.shape)}"
+            )
+        timestep_indices = timestep_indices.to(hidden_states.device)
+        if not torch.all((timestep_indices >= 0) & (timestep_indices < shift.shape[0])):
+            raise ValueError(f"MiniMax-H3 final {label} timestep indices are out of range")
+        row_shift = shift[timestep_indices]
+        row_scale = scale[timestep_indices]
+        return hidden_states * (1.0 + row_scale) + row_shift
+
+
+@dataclass(frozen=True)
+class MiniMaxH3Output:
+    video: torch.Tensor
+    audio: torch.Tensor
+
+
+class MiniMaxH3Model(nn.Module):
+    def __init__(
+        self,
+        config: MiniMaxH3Config,
+        *,
+        attn_mode: str = "torch",
+        split_attn: bool = False,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        if attn_mode not in {"torch", "sdpa", "flash", "flash3", "sageattn", "xformers"}:
+            raise ValueError(f"Unsupported MiniMax-H3 attention mode: {attn_mode}")
+        self.config = config
+        self.attn_mode = attn_mode
+        self.split_attn = split_attn
+        self.video_patch_proj = nn.Linear(config.video_patch_dim, config.hidden_size, dtype=torch.float32, device=device)
+        self.audio_patch_proj = nn.Linear(config.audio_in_channels, config.hidden_size, dtype=torch.float32, device=device)
+        self.condition_proj = nn.Linear(config.text_dim, config.hidden_size, dtype=dtype, device=device)
+        self.time_embedder = TimeEmbedder(
+            config.timestep_input_dim,
+            config.time_embed_hidden_size,
+            config.time_embed_dim,
+            device=device,
+        )
+        self.rope = nn.Module()
+        inv_freq = torch.empty(config.rope_inv_freq_len, dtype=torch.float32, device=device)
+        self.rope.register_buffer("inv_freq", inv_freq)
+        self.token_refiner = TokenRefiner(
+            config,
+            attn_mode=attn_mode,
+            split_attn=split_attn,
+            dtype=dtype,
+            device=device,
+        )
+        self.blocks = nn.ModuleList(
+            [
+                DiTBlock(config, attn_mode=attn_mode, split_attn=split_attn, dtype=dtype, device=device)
+                for _ in range(config.num_layers)
+            ]
+        )
+        self.final_layer = FinalLayer(
+            config.hidden_size,
+            config.time_embed_dim,
+            config.video_patch_dim,
+            config.audio_in_channels,
+            norm_eps=config.final_norm_eps,
+            dtype=dtype,
+            device=device,
+        )
+
+        self.gradient_checkpointing = False
+        self.activation_cpu_offloading = False
+        self.blocks_to_swap = 0
+        self.offloader = None
+        self._execution_device = torch.device(device) if device is not None else None
+        self._rotary_cache: OrderedDict[tuple[H3PackedLayout, torch.device, torch.dtype], torch.Tensor] = OrderedDict()
+
+    def _apply(self, fn, recurse: bool = True):
+        result = super()._apply(fn, recurse=recurse)
+        if hasattr(self, "_rotary_cache"):
+            self._rotary_cache.clear()
+        return result
+
+    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
+        self._rotary_cache.clear()
+        return super().load_state_dict(state_dict, strict=strict, assign=assign)
+
+    @property
+    def device(self) -> torch.device:
+        return self.video_patch_proj.weight.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.condition_proj.weight.dtype
+
+    def enable_gradient_checkpointing(self, activation_cpu_offloading: bool = False) -> None:
+        self.gradient_checkpointing = True
+        self.activation_cpu_offloading = activation_cpu_offloading
+
+    def disable_gradient_checkpointing(self) -> None:
+        self.gradient_checkpointing = False
+        self.activation_cpu_offloading = False
+
+    def enable_block_swap(self, num_blocks: int, config: BlockSwapConfig) -> None:
+        if num_blocks <= 0:
+            raise ValueError("MiniMax-H3 blocks to swap must be positive")
+        if num_blocks > len(self.blocks) - 2:
+            raise ValueError(f"MiniMax-H3 cannot swap more than {len(self.blocks) - 2} of {len(self.blocks)} blocks")
+        self.blocks_to_swap = num_blocks
+        self._execution_device = torch.device(config.device)
+        self.offloader = create_offloader(
+            "minimax-h3",
+            self.blocks,
+            len(self.blocks),
+            num_blocks,
+            config,
+        )
+
+    def move_to_device_except_swap_blocks(self, device: torch.device) -> None:
+        device = torch.device(device)
+        if self.blocks_to_swap:
+            saved_blocks = self.blocks
+            self.blocks = nn.ModuleList()
+        self.to(device)
+        if self.blocks_to_swap:
+            self.blocks = saved_blocks
+        self._execution_device = device
+
+    def prepare_block_swap_before_forward(self) -> None:
+        if not self.blocks_to_swap:
+            return
+        if self.offloader is None:
+            raise RuntimeError("MiniMax-H3 block swap has no offloader")
+        self.offloader.prepare_block_devices_before_forward(self.blocks)
+
+    def switch_block_swap_for_inference(self) -> None:
+        if self.blocks_to_swap:
+            self.offloader.set_forward_only(True)
+            self.prepare_block_swap_before_forward()
+
+    def switch_block_swap_for_training(self) -> None:
+        if self.blocks_to_swap:
+            self.offloader.set_forward_only(False)
+            self.prepare_block_swap_before_forward()
+
+    def _assert_block_device(self, block: nn.Module, index: int) -> None:
+        expected = self._execution_device
+        if expected is None:
+            return
+
+        def matches(actual: torch.device) -> bool:
+            return actual.type == expected.type and (expected.index is None or actual.index == expected.index)
+
+        for name, parameter in block.named_parameters():
+            if not matches(parameter.device):
+                raise RuntimeError(
+                    f"MiniMax-H3 block {index} parameter {name} is on {parameter.device}, expected {expected} after wait"
+                )
+        for name, buffer in block.named_buffers():
+            if not matches(buffer.device):
+                raise RuntimeError(f"MiniMax-H3 block {index} buffer {name} is on {buffer.device}, expected {expected} after wait")
+
+    def _rotation_table(self, position_ids: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        positions = position_ids.to(device=self.rope.inv_freq.device, dtype=torch.float32)
+        if positions.ndim != 3 or positions.shape[-1] != 3:
+            raise ValueError("MiniMax-H3 position ids must be [B,S,3]")
+        per_axis = positions.unsqueeze(-1) * self.rope.inv_freq.reshape(1, 1, 1, -1)
+        temporal, height, width = per_axis.unbind(dim=2)
+        half = torch.cat((temporal, height, width), dim=-1)
+        angles = torch.cat((half, half), dim=-1)
+        pair_count = angles.shape[-1] // 2
+        angles = angles[..., :pair_count]
+        cosine, sine = torch.cos(angles), torch.sin(angles)
+        return (
+            torch.stack((cosine, -sine, sine, cosine), dim=-1)
+            .reshape(
+                angles.shape[0],
+                angles.shape[1],
+                1,
+                pair_count,
+                2,
+                2,
+            )
+            .to(dtype)
+        )
+
+    def _cached_rotation_table(
+        self,
+        layout: H3PackedLayout,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        key = (layout, torch.device(device), dtype)
+        cached = self._rotary_cache.get(key)
+        if cached is not None:
+            self._rotary_cache.move_to_end(key)
+            return cached
+        position_ids = build_position_grid(layout, device=device)
+        rotation_table = self._rotation_table(position_ids, dtype)
+        self._rotary_cache[key] = rotation_table
+        while len(self._rotary_cache) > _ROTARY_CACHE_SIZE:
+            self._rotary_cache.popitem(last=False)
+        return rotation_table
+
+    @staticmethod
+    def _validate_text_tags(tags: torch.Tensor, batch_size: int, text_length: int) -> torch.Tensor:
+        tags = torch.as_tensor(tags)
+        if tags.shape != (batch_size, text_length):
+            raise ValueError(f"MiniMax-H3 text token tags must be [{batch_size},{text_length}]")
+        return tags
+
+    @staticmethod
+    def _validate_condition_rows(
+        latents: Sequence[torch.Tensor],
+        segments,
+        pack,
+        *,
+        batch_size: int,
+        label: str,
+    ) -> list[torch.Tensor]:
+        if len(latents) != len(segments):
+            raise ValueError(f"MiniMax-H3 expected {len(segments)} {label} tensors, got {len(latents)}")
+        rows = []
+        for tensor, segment in zip(latents, segments):
+            if tensor.shape[0] != batch_size:
+                raise ValueError(f"MiniMax-H3 {segment.role} batch size does not match the targets")
+            packed = pack(tensor)
+            if packed.shape[1] != segment.row_count:
+                raise ValueError(f"MiniMax-H3 {segment.role} has {packed.shape[1]} packed rows, expected {segment.row_count}")
+            rows.append(packed)
+        return rows
+
+    @staticmethod
+    def _target_denoise_rows(
+        mask: torch.Tensor | None,
+        latents: torch.Tensor,
+        pack,
+        *,
+        label: str,
+    ) -> torch.Tensor | None:
+        if mask is None:
+            return None
+        mask = torch.as_tensor(mask, device=latents.device)
+        if mask.is_complex():
+            raise ValueError(f"MiniMax-H3 {label} must be real")
+        if mask.is_floating_point() and not torch.isfinite(mask).all():
+            raise ValueError(f"MiniMax-H3 {label} must be finite")
+        if not torch.all((mask == 0) | (mask == 1)):
+            raise ValueError(f"MiniMax-H3 {label} must contain only 0 (fixed) and 1 (regenerate)")
+        try:
+            mask = torch.broadcast_to(mask, latents.shape)
+        except RuntimeError as error:
+            raise ValueError(
+                f"MiniMax-H3 {label} shape {tuple(mask.shape)} is not broadcastable to {tuple(latents.shape)}"
+            ) from error
+        packed = pack(mask.to(torch.bool))
+        row_values = packed[..., :1]
+        if not torch.equal(packed, row_values.expand_as(packed)):
+            raise ValueError(f"MiniMax-H3 {label} must be constant within every packed target row")
+        return row_values[..., 0].detach().cpu()
+
+    def forward(
+        self,
+        *,
+        video_latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        text_token_tags: torch.Tensor,
+        layout: H3PackedLayout,
+        model_t_video: float | torch.Tensor,
+        model_t_audio: float | torch.Tensor,
+        visual_condition_latents: Sequence[torch.Tensor] = (),
+        audio_condition_latents: Sequence[torch.Tensor] = (),
+        visual_condition_clean: float = 0.999,
+        audio_condition_clean: float = 1.0,
+        video_denoise_mask: torch.Tensor | None = None,
+        audio_denoise_mask: torch.Tensor | None = None,
+    ) -> MiniMaxH3Output:
+        if video_latents.ndim != 5 or tuple(video_latents.shape[2:]) != (
+            layout.target_video.frames,
+            layout.target_video.height,
+            layout.target_video.width,
+        ):
+            raise ValueError("MiniMax-H3 target video tensor does not match the packed layout")
+        if audio_latents.ndim != 4 or tuple(audio_latents.shape[1:]) != (
+            self.config.audio_in_channels,
+            2,
+            layout.target_audio_frames,
+        ):
+            raise ValueError("MiniMax-H3 target audio tensor does not match the packed layout")
+        batch_size = video_latents.shape[0]
+        if batch_size != 1:
+            raise ValueError(f"MiniMax-H3 R1 requires batch_size=1, got {batch_size}; use gradient accumulation")
+        if audio_latents.shape[0] != batch_size:
+            raise ValueError("MiniMax-H3 target video and audio batch sizes differ")
+        if text_hidden_states.shape != (batch_size, layout.text_length, self.config.text_dim):
+            raise ValueError(f"MiniMax-H3 text hidden states must be [{batch_size},{layout.text_length},{self.config.text_dim}]")
+        text_token_tags = self._validate_text_tags(text_token_tags, batch_size, layout.text_length)
+
+        execution_device = self.video_patch_proj.weight.device
+        video_dtype = video_latents.dtype
+        audio_dtype = audio_latents.dtype
+        video_latents = video_latents.to(execution_device)
+        audio_latents = audio_latents.to(execution_device)
+        text_hidden_states = text_hidden_states.to(execution_device)
+        visual_condition_latents = tuple(tensor.to(execution_device) for tensor in visual_condition_latents)
+        audio_condition_latents = tuple(tensor.to(execution_device) for tensor in audio_condition_latents)
+        video_denoise_rows = self._target_denoise_rows(
+            video_denoise_mask,
+            video_latents,
+            pack_video_rows,
+            label="video denoise mask",
+        )
+        audio_denoise_rows = self._target_denoise_rows(
+            audio_denoise_mask,
+            audio_latents,
+            pack_audio_rows,
+            label="audio denoise mask",
+        )
+
+        visual_segments = tuple(segment for segment in layout.segments if segment.kind == "visual_condition")
+        audio_segments = tuple(segment for segment in layout.segments if segment.kind == "audio_condition")
+        if layout.task == "fl2va":
+            expected_visual_geometries = layout.visual_conditions
+        elif layout.task == "ref2va":
+            expected_visual_geometries = tuple(
+                reference.video for reference in layout.references if reference.kind in {"image", "video"}
+            )
+        else:
+            expected_visual_geometries = ()
+        expected_audio_frames = (
+            tuple(reference.audio_frames for reference in layout.references if reference.audio_frames)
+            if layout.task == "ref2va"
+            else ()
+        )
+        for tensor, segment, geometry in zip(
+            visual_condition_latents,
+            visual_segments,
+            expected_visual_geometries,
+        ):
+            actual = tuple(tensor.shape[2:]) if tensor.ndim == 5 else tuple(tensor.shape)
+            expected = (geometry.frames, geometry.height, geometry.width)
+            if actual != expected:
+                raise ValueError(
+                    f"MiniMax-H3 {segment.role} geometry expected {'x'.join(map(str, expected))}, got {'x'.join(map(str, actual))}"
+                )
+        for tensor, segment, frames in zip(audio_condition_latents, audio_segments, expected_audio_frames):
+            actual = tuple(tensor.shape[1:]) if tensor.ndim == 4 else tuple(tensor.shape)
+            expected = (self.config.audio_in_channels, 2, frames)
+            if actual != expected:
+                raise ValueError(
+                    f"MiniMax-H3 {segment.role} geometry expected {'x'.join(map(str, expected))}, got {'x'.join(map(str, actual))}"
+                )
+        visual_rows = self._validate_condition_rows(
+            visual_condition_latents,
+            visual_segments,
+            pack_video_rows,
+            batch_size=batch_size,
+            label="visual condition",
+        )
+        audio_rows = self._validate_condition_rows(
+            audio_condition_latents,
+            audio_segments,
+            pack_audio_rows,
+            batch_size=batch_size,
+            label="audio condition",
+        )
+        visual_rows.append(pack_video_rows(video_latents))
+        audio_rows.append(pack_audio_rows(audio_latents))
+        projected_video = self.video_patch_proj(torch.cat(visual_rows, dim=1).to(torch.float32)).to(self.dtype)
+        projected_audio = self.audio_patch_proj(torch.cat(audio_rows, dim=1).to(torch.float32)).to(self.dtype)
+
+        text = self.condition_proj(text_hidden_states.to(self.dtype))
+        text = self.token_refiner(text)
+        parts = []
+        visual_offset = 0
+        audio_offset = 0
+        for segment in layout.segments:
+            if segment.kind == "text":
+                parts.append(text)
+            elif segment.kind in {"visual_condition", "target_video"}:
+                parts.append(projected_video[:, visual_offset : visual_offset + segment.row_count])
+                visual_offset += segment.row_count
+            else:
+                parts.append(projected_audio[:, audio_offset : audio_offset + segment.row_count])
+                audio_offset += segment.row_count
+        hidden_states = torch.cat(parts, dim=1)
+        if hidden_states.shape[1] != layout.row_count:
+            raise RuntimeError("MiniMax-H3 packed residual stream length drifted from its layout")
+
+        timestep_rows = build_timestep_rows(
+            layout,
+            text_token_tags=text_token_tags,
+            model_t_video=model_t_video,
+            model_t_audio=model_t_audio,
+            visual_condition_clean=visual_condition_clean,
+            audio_condition_clean=audio_condition_clean,
+            target_video_denoise_rows=video_denoise_rows,
+            target_audio_denoise_rows=audio_denoise_rows,
+        )
+        timestep_embeddings = self.time_embedder(timestep_rows.unique_timesteps.to(execution_device)).to(self.dtype)
+        rotation_table = self._cached_rotation_table(
+            layout,
+            device=execution_device,
+            dtype=hidden_states.dtype,
+        )
+
+        for index, block in enumerate(self.blocks):
+            if self.blocks_to_swap:
+                self.offloader.wait_for_block(index)
+                self._assert_block_device(block, index)
+            if self.gradient_checkpointing and self.training:
+                forward_fn = block
+                if self.activation_cpu_offloading:
+                    forward_fn = create_cpu_offloading_wrapper(forward_fn, execution_device)
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    forward_fn,
+                    hidden_states,
+                    timestep_embeddings,
+                    timestep_rows.block_segments,
+                    rotation_table,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    timestep_embeddings,
+                    timestep_rows.block_segments,
+                    rotation_table,
+                )
+            if self.blocks_to_swap:
+                self.offloader.submit_move_blocks_forward(self.blocks, index)
+
+        final_layer_kwargs = {
+            "video_slice": layout.target_video_segment.row_slice,
+            "audio_slice": layout.target_audio_segment.row_slice,
+            "video_timestep_index": timestep_rows.video_timestep_index,
+            "audio_timestep_index": timestep_rows.audio_timestep_index,
+        }
+        if video_denoise_rows is not None or audio_denoise_rows is not None:
+            final_layer_kwargs.update(
+                video_timestep_indices=timestep_rows.video_timestep_indices,
+                audio_timestep_indices=timestep_rows.audio_timestep_indices,
+            )
+        video_rows, audio_rows = self.final_layer(hidden_states, timestep_embeddings, **final_layer_kwargs)
+        video, audio = unpack_targets(layout, video_rows, audio_rows)
+        return MiniMaxH3Output(video=video.to(video_dtype), audio=audio.to(audio_dtype))
+
+
+def load_h3_transformer(
+    checkpoint_path: str | Path,
+    *,
+    device: torch.device | str,
+    dtype: torch.dtype = torch.bfloat16,
+    attn_mode: str = "torch",
+    split_attn: bool = False,
+    disable_mmap: bool = False,
+) -> MiniMaxH3Model:
+    if dtype != torch.bfloat16:
+        raise ValueError("MiniMax-H3 R1 accepts only BF16 transformer checkpoints")
+    files = resolve_safetensors_files(checkpoint_path)
+    config = parse_h3_transformer_config(load_safetensors_metadata(files))
+    return load_safetensors_module(
+        lambda: MiniMaxH3Model(config, attn_mode=attn_mode, split_attn=split_attn, dtype=dtype),
+        files,
+        device=device,
+        dtype=None,
+        strict_dtype=True,
+        disable_mmap=disable_mmap,
+    )
diff --git a/musubi_tuner/minimax_h3/packing.py b/musubi_tuner/minimax_h3/packing.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1d60d9cad572b224d34e3cc77622d7af17b0588
--- /dev/null
+++ b/musubi_tuner/minimax_h3/packing.py
@@ -0,0 +1,591 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted for Musubi from Hugging Face Diffusers PR #14355 at commit
+# abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc
+# (modular_pipelines/minimax_h3/packing.py and packing_ref2va.py).
+# ComfyUI is used only as an independent numerical reference.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+
+from musubi_tuner.minimax_h3.media import H3Task
+
+VIDEO_CHANNELS = 24
+AUDIO_CHANNELS = 32
+STEREO_CHANNELS = 2
+VIDEO_PATCH_SIZE = (1, 2, 2)
+FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
+FRAME_RESCALE = 5.0 / 3.0
+
+H3ReferenceKind = Literal["image", "audio", "video"]
+H3SegmentKind = Literal["text", "visual_condition", "audio_condition", "target_audio", "target_video"]
+
+
+@dataclass(frozen=True)
+class H3VideoGeometry:
+    frames: int
+    height: int
+    width: int
+
+    def __post_init__(self) -> None:
+        if self.frames <= 0:
+            raise ValueError(f"MiniMax-H3 video geometry needs positive frames, got {self.frames}")
+        if self.height <= 0 or self.width <= 0 or self.height % 2 or self.width % 2:
+            raise ValueError(f"MiniMax-H3 latent video axes must be positive multiples of 2, got {self.height}x{self.width}")
+
+    @property
+    def frame_rows(self) -> int:
+        return (self.height // 2) * (self.width // 2)
+
+    @property
+    def row_count(self) -> int:
+        return self.frames * self.frame_rows
+
+
+@dataclass(frozen=True)
+class H3ReferenceGeometry:
+    kind: H3ReferenceKind
+    video: H3VideoGeometry | None = None
+    audio_frames: int = 0
+
+    def __post_init__(self) -> None:
+        if self.kind not in {"image", "audio", "video"}:
+            raise ValueError(f"Unsupported MiniMax-H3 reference kind: {self.kind}")
+        if self.audio_frames < 0:
+            raise ValueError(f"MiniMax-H3 reference audio length must be nonnegative, got {self.audio_frames}")
+        if self.kind == "image":
+            if self.video is None or self.video.frames != 1 or self.audio_frames:
+                raise ValueError("MiniMax-H3 image references require one visual frame and no audio rows")
+        elif self.kind == "audio":
+            if self.video is not None or self.audio_frames <= 0:
+                raise ValueError("MiniMax-H3 audio references require audio rows and no visual geometry")
+        else:
+            if self.video is None:
+                raise ValueError("MiniMax-H3 video references require visual geometry")
+            _validate_video_latent_frames(self.video.frames, "reference video")
+            if self.audio_frames:
+                expected = _expected_audio_frames(self.video.frames)
+                if self.audio_frames != expected:
+                    raise ValueError(f"MiniMax-H3 reference video audio has {self.audio_frames} frames, expected {expected}")
+
+
+@dataclass(frozen=True)
+class H3RowSegment:
+    role: str
+    kind: H3SegmentKind
+    start: int
+    stop: int
+
+    def __post_init__(self) -> None:
+        if self.start < 0 or self.stop <= self.start:
+            raise ValueError(f"Invalid MiniMax-H3 row segment {self.role}: [{self.start}, {self.stop})")
+
+    @property
+    def row_count(self) -> int:
+        return self.stop - self.start
+
+    @property
+    def row_slice(self) -> slice:
+        return slice(self.start, self.stop)
+
+
+@dataclass(frozen=True)
+class H3PackedLayout:
+    task: H3Task
+    text_length: int
+    target_video: H3VideoGeometry
+    target_audio_frames: int
+    visual_conditions: tuple[H3VideoGeometry, ...]
+    references: tuple[H3ReferenceGeometry, ...]
+    segments: tuple[H3RowSegment, ...]
+    row_count: int
+
+    def segment(self, role: str) -> H3RowSegment:
+        matches = [segment for segment in self.segments if segment.role == role]
+        if len(matches) != 1:
+            raise KeyError(f"MiniMax-H3 layout expected exactly one {role!r} segment, found {len(matches)}")
+        return matches[0]
+
+    @property
+    def target_video_segment(self) -> H3RowSegment:
+        return self.segment("target_video")
+
+    @property
+    def target_audio_segment(self) -> H3RowSegment:
+        return self.segment("target_audio")
+
+
+@dataclass(frozen=True)
+class H3TimestepRows:
+    unique_timesteps: torch.Tensor
+    row_timesteps: torch.Tensor
+    row_timestep_indices: torch.Tensor
+    token_tags: torch.Tensor
+    block_adaln_indices: torch.Tensor
+    block_segments: tuple[tuple[tuple[int, int, int], ...], ...]
+    video_timestep_index: int
+    audio_timestep_index: int
+    video_timestep_indices: torch.Tensor
+    audio_timestep_indices: torch.Tensor
+
+
+def _coerce_video_geometry(value: H3VideoGeometry | Sequence[int], label: str) -> H3VideoGeometry:
+    if isinstance(value, H3VideoGeometry):
+        return value
+    if len(value) != 3:
+        raise ValueError(f"MiniMax-H3 {label} geometry must be (frames,height,width), got {tuple(value)}")
+    return H3VideoGeometry(*(int(part) for part in value))
+
+
+def _validate_video_latent_frames(frames: int, label: str) -> None:
+    if frames < 2 or (frames - 2) % 5:
+        raise ValueError(f"MiniMax-H3 {label} latent frames must be 5*n+2, got {frames}")
+
+
+def _expected_audio_frames(video_latent_frames: int) -> int:
+    _validate_video_latent_frames(video_latent_frames, "target video")
+    n = (video_latent_frames - 2) // 5
+    pixel_frames = 17 * n + 5
+    return (10 * pixel_frames + 3) // 6
+
+
+def pack_video_rows(latents: torch.Tensor) -> torch.Tensor:
+    """Patch H3 video latents while preserving the leading batch axis."""
+    unbatched = latents.ndim == 4
+    if unbatched:
+        latents = latents.unsqueeze(0)
+    if latents.ndim != 5:
+        raise ValueError(f"MiniMax-H3 video latents must be [B,24,F,H,W], got {tuple(latents.shape)}")
+    batch, channels, frames, height, width = latents.shape
+    if channels != VIDEO_CHANNELS:
+        raise ValueError(f"MiniMax-H3 video latents need {VIDEO_CHANNELS} channels, got {channels}")
+    patch_t, patch_h, patch_w = VIDEO_PATCH_SIZE
+    if frames % patch_t or height % patch_h or width % patch_w:
+        raise ValueError(f"MiniMax-H3 video latent shape {frames}x{height}x{width} is not divisible by patch {VIDEO_PATCH_SIZE}")
+    rows = latents.reshape(
+        batch,
+        channels,
+        frames // patch_t,
+        patch_t,
+        height // patch_h,
+        patch_h,
+        width // patch_w,
+        patch_w,
+    )
+    rows = rows.permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(
+        batch,
+        (frames // patch_t) * (height // patch_h) * (width // patch_w),
+        channels * patch_t * patch_h * patch_w,
+    )
+    return rows[0] if unbatched else rows
+
+
+def pack_audio_rows(latents: torch.Tensor) -> torch.Tensor:
+    """Pack stereo H3 audio in channel-major order."""
+    unbatched = latents.ndim == 3
+    if unbatched:
+        latents = latents.unsqueeze(0)
+    if latents.ndim != 4 or latents.shape[1] != AUDIO_CHANNELS or latents.shape[2] != STEREO_CHANNELS:
+        raise ValueError(f"MiniMax-H3 audio latents must be [B,32,2,A], got {tuple(latents.shape)}")
+    batch, channels, stereo, frames = latents.shape
+    rows = latents.permute(0, 2, 3, 1).reshape(batch, stereo * frames, channels)
+    return rows[0] if unbatched else rows
+
+
+def build_h3_layout(
+    *,
+    task: H3Task,
+    text_length: int,
+    target_video: H3VideoGeometry | Sequence[int],
+    target_audio_frames: int,
+    visual_conditions: Sequence[H3VideoGeometry | Sequence[int]] = (),
+    references: Sequence[H3ReferenceGeometry] = (),
+) -> H3PackedLayout:
+    if task not in {"t2va", "fl2va", "ref2va"}:
+        raise ValueError(f"Unsupported MiniMax-H3 task: {task}")
+    if text_length <= 0:
+        raise ValueError(f"MiniMax-H3 text length must be positive, got {text_length}")
+    target_video = _coerce_video_geometry(target_video, "target video")
+    _validate_video_latent_frames(target_video.frames, "target video")
+    expected_audio_frames = _expected_audio_frames(target_video.frames)
+    if target_audio_frames != expected_audio_frames:
+        raise ValueError(
+            f"MiniMax-H3 target audio has {target_audio_frames} frames, expected {expected_audio_frames} "
+            f"for {target_video.frames} video latent frames"
+        )
+
+    visual_conditions = tuple(
+        _coerce_video_geometry(condition, f"visual condition {index}") for index, condition in enumerate(visual_conditions)
+    )
+    references = tuple(references)
+    if task == "t2va" and (visual_conditions or references):
+        raise ValueError("MiniMax-H3 T2VA layout does not accept condition rows")
+    if task == "fl2va":
+        if references or len(visual_conditions) != 2:
+            raise ValueError("MiniMax-H3 FL2VA layout requires exactly first and last visual conditions")
+        for role, condition in zip(("first", "last"), visual_conditions):
+            if condition != H3VideoGeometry(1, target_video.height, target_video.width):
+                raise ValueError(f"MiniMax-H3 FL2VA {role} condition must be one target-sized latent frame, got {condition}")
+    if task == "ref2va":
+        if visual_conditions or not references:
+            raise ValueError("MiniMax-H3 Ref2VA layout requires ordered references and no FL2VA conditions")
+        if not any(reference.kind in {"image", "video"} for reference in references):
+            raise ValueError("MiniMax-H3 Ref2VA layout requires at least one visual reference")
+
+    segments = []
+    row = 0
+
+    def append(role: str, kind: H3SegmentKind, row_count: int) -> None:
+        nonlocal row
+        segments.append(H3RowSegment(role, kind, row, row + row_count))
+        row += row_count
+
+    append("text", "text", text_length)
+    if task == "fl2va":
+        append("first", "visual_condition", visual_conditions[0].row_count)
+        append("last", "visual_condition", visual_conditions[1].row_count)
+    elif task == "ref2va":
+        for index, reference in enumerate(references):
+            prefix = f"ref_{index:03d}"
+            if reference.kind == "image":
+                append(f"{prefix}_image", "visual_condition", reference.video.row_count)
+            elif reference.kind == "audio":
+                append(f"{prefix}_audio", "audio_condition", reference.audio_frames * STEREO_CHANNELS)
+            else:
+                if reference.audio_frames:
+                    append(f"{prefix}_audio", "audio_condition", reference.audio_frames * STEREO_CHANNELS)
+                append(f"{prefix}_video", "visual_condition", reference.video.row_count)
+    append("target_audio", "target_audio", target_audio_frames * STEREO_CHANNELS)
+    append("target_video", "target_video", target_video.row_count)
+
+    return H3PackedLayout(
+        task=task,
+        text_length=text_length,
+        target_video=target_video,
+        target_audio_frames=target_audio_frames,
+        visual_conditions=visual_conditions,
+        references=references,
+        segments=tuple(segments),
+        row_count=row,
+    )
+
+
+def _axis_from_sqrt_area(dimension: int, sqrt_area: float) -> torch.Tensor:
+    ratio = dimension / sqrt_area
+    count = dimension // 2
+    return (torch.arange(count, dtype=torch.float64) * (ratio / count) + (1.0 - ratio) / 2.0) * 32.0
+
+
+def _frame_grid(video: H3VideoGeometry) -> tuple[torch.Tensor, torch.Tensor]:
+    sqrt_area = math.sqrt(video.height * video.width)
+    height_axis = _axis_from_sqrt_area(video.height, sqrt_area)
+    width_axis = _axis_from_sqrt_area(video.width, sqrt_area)
+    height, width = torch.meshgrid(height_axis, width_axis, indexing="ij")
+    return torch.stack((height.reshape(-1), width.reshape(-1)), dim=-1), width_axis
+
+
+def _video_t_spans(frames: int) -> list[float]:
+    return [FRAME_RESCALE * FRAME_PER_TOKEN[index % len(FRAME_PER_TOKEN)] for index in range(frames)]
+
+
+def _video_grid(video: H3VideoGeometry, cursor: float) -> torch.Tensor:
+    spans = torch.tensor(_video_t_spans(video.frames), dtype=torch.float64)
+    times = cursor + torch.cat((torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)))
+    frame, _ = _frame_grid(video)
+    grid = torch.empty(video.frames, frame.shape[0], 3, dtype=torch.float64)
+    grid[..., 0] = times[:, None]
+    grid[..., 1:] = frame[None]
+    return grid.reshape(-1, 3)
+
+
+def _audio_grid(cursor: float, frames: int, width_low: float, width_high: float) -> torch.Tensor:
+    grid = torch.zeros(frames * STEREO_CHANNELS, 3, dtype=torch.float64)
+    grid[:, 0] = (cursor + torch.arange(frames, dtype=torch.float64)).repeat(STEREO_CHANNELS)
+    grid[:frames, 2] = width_low
+    grid[frames:, 2] = width_high
+    return grid
+
+
+def build_position_grid(layout: H3PackedLayout, *, device: torch.device | str | None = None) -> torch.Tensor:
+    positions = torch.empty(layout.row_count, 3, dtype=torch.float64)
+    text = layout.segment("text")
+    positions[text.row_slice] = 0
+    positions[text.row_slice, 0] = torch.arange(layout.text_length, dtype=torch.float64)
+
+    target_frame, target_width = _frame_grid(layout.target_video)
+    target_width_endpoints = float(target_width[0]), float(target_width[-1])
+    cursor = float(layout.text_length)
+
+    if layout.task == "fl2va":
+        first = layout.segment("first")
+        positions[first.row_slice, 0] = cursor
+        positions[first.row_slice, 1:] = target_frame
+        last = layout.segment("last")
+        last_time = cursor + sum(_video_t_spans(layout.target_video.frames)) - FRAME_RESCALE
+        positions[last.row_slice, 0] = last_time
+        positions[last.row_slice, 1:] = target_frame
+    elif layout.task == "ref2va":
+        for index, reference in enumerate(layout.references):
+            prefix = f"ref_{index:03d}"
+            if reference.kind == "image":
+                segment = layout.segment(f"{prefix}_image")
+                frame, _ = _frame_grid(reference.video)
+                positions[segment.row_slice, 0] = cursor
+                positions[segment.row_slice, 1:] = frame
+                cursor += 1.0
+            elif reference.kind == "audio":
+                segment = layout.segment(f"{prefix}_audio")
+                positions[segment.row_slice] = _audio_grid(cursor, reference.audio_frames, *target_width_endpoints)
+                cursor += float(reference.audio_frames)
+            else:
+                frame, width = _frame_grid(reference.video)
+                if reference.audio_frames:
+                    audio = layout.segment(f"{prefix}_audio")
+                    positions[audio.row_slice] = _audio_grid(
+                        cursor,
+                        reference.audio_frames,
+                        float(width[0]),
+                        float(width[-1]),
+                    )
+                video = layout.segment(f"{prefix}_video")
+                positions[video.row_slice] = _video_grid(reference.video, cursor)
+                cursor += max(float(reference.audio_frames), sum(_video_t_spans(reference.video.frames)))
+
+    target_audio = layout.target_audio_segment
+    positions[target_audio.row_slice] = _audio_grid(
+        cursor,
+        layout.target_audio_frames,
+        *target_width_endpoints,
+    )
+    target_video = layout.target_video_segment
+    positions[target_video.row_slice] = _video_grid(layout.target_video, cursor)
+    positions = positions.unsqueeze(0)
+    return positions.to(device=device) if device is not None else positions
+
+
+def _single_model_time(value: float | torch.Tensor, label: str) -> float:
+    tensor = torch.as_tensor(value).detach().reshape(-1)
+    if tensor.numel() != 1:
+        raise ValueError(f"MiniMax-H3 R1 {label} must contain exactly one value")
+    scalar = float(tensor[0].item())
+    if not math.isfinite(scalar) or not 0.0 <= scalar <= 1.0:
+        raise ValueError(f"MiniMax-H3 {label} must be finite and in [0,1], got {scalar}")
+    return scalar
+
+
+def _validate_clean_coefficient(value: float, label: str) -> float:
+    value = float(value)
+    if not math.isfinite(value) or not 0.0 <= value <= 1.0:
+        raise ValueError(f"MiniMax-H3 {label} must be finite and in [0,1], got {value}")
+    return value
+
+
+def _build_modulation_segments(indices: torch.Tensor) -> tuple[tuple[int, int, int], ...]:
+    change_points = torch.nonzero(indices[1:] != indices[:-1], as_tuple=False).flatten() + 1
+    boundaries = torch.cat((indices.new_tensor([0]), change_points, indices.new_tensor([indices.numel()])))
+    return tuple(
+        tuple(row)
+        for row in torch.stack(
+            (boundaries[:-1], boundaries[1:], indices[boundaries[:-1]]),
+            dim=1,
+        ).tolist()
+    )
+
+
+def _validate_target_denoise_rows(mask: torch.Tensor | None, row_count: int, label: str) -> torch.Tensor:
+    if mask is None:
+        return torch.ones((1, row_count), dtype=torch.bool)
+    mask = torch.as_tensor(mask).detach().cpu()
+    if mask.shape != (1, row_count):
+        raise ValueError(f"MiniMax-H3 {label} must be [1,{row_count}], got {tuple(mask.shape)}")
+    if mask.is_floating_point() and not torch.isfinite(mask).all():
+        raise ValueError(f"MiniMax-H3 {label} must be finite")
+    if not torch.all((mask == 0) | (mask == 1)):
+        raise ValueError(f"MiniMax-H3 {label} must contain only 0 (fixed) and 1 (regenerate)")
+    return mask.to(torch.bool)
+
+
+def _timestep_index(unique_timesteps: torch.Tensor, value: float, label: str, *, fallback: int) -> int:
+    matches = torch.nonzero(unique_timesteps == torch.tensor(value, dtype=unique_timesteps.dtype), as_tuple=False)
+    if matches.numel() == 0:
+        return fallback
+    if matches.shape != (1, 1):
+        raise RuntimeError(f"MiniMax-H3 {label} timestep is missing from the packed plan")
+    return int(matches[0, 0])
+
+
+def build_timestep_rows(
+    layout: H3PackedLayout,
+    *,
+    text_token_tags: torch.Tensor,
+    model_t_video: float | torch.Tensor,
+    model_t_audio: float | torch.Tensor,
+    visual_condition_clean: float = 0.999,
+    audio_condition_clean: float = 1.0,
+    target_video_denoise_rows: torch.Tensor | None = None,
+    target_audio_denoise_rows: torch.Tensor | None = None,
+) -> H3TimestepRows:
+    video_time = _single_model_time(model_t_video, "video model time")
+    audio_time = _single_model_time(model_t_audio, "audio model time")
+    visual_clean = _validate_clean_coefficient(visual_condition_clean, "visual condition clean coefficient")
+    audio_clean = _validate_clean_coefficient(audio_condition_clean, "audio condition clean coefficient")
+
+    text_token_tags = torch.as_tensor(text_token_tags)
+    if text_token_tags.dtype != torch.int64 or text_token_tags.shape != (1, layout.text_length):
+        raise ValueError(f"MiniMax-H3 text token tags must be int64 [1,{layout.text_length}]")
+    if not torch.all((text_token_tags == 0) | (text_token_tags == 1)):
+        raise ValueError("MiniMax-H3 text token tags may contain only 0 and 1")
+    text_token_tags = text_token_tags.detach().cpu()
+
+    segment_times = {}
+    for segment in layout.segments:
+        if segment.kind in {"text", "target_video"}:
+            segment_times[segment.role] = video_time
+        elif segment.kind == "target_audio":
+            segment_times[segment.role] = audio_time
+        elif segment.kind == "visual_condition":
+            segment_times[segment.role] = max(video_time, visual_clean)
+        else:
+            segment_times[segment.role] = max(audio_time, audio_clean)
+
+    row_timesteps = torch.empty((1, layout.row_count), dtype=torch.float32)
+    token_tags = torch.empty((1, layout.row_count), dtype=torch.int64)
+    for segment in layout.segments:
+        value = segment_times[segment.role]
+        row_timesteps[:, segment.row_slice] = value
+        if segment.kind == "text":
+            token_tags[:, segment.row_slice] = text_token_tags
+        elif segment.kind in {"visual_condition", "target_video"}:
+            token_tags[:, segment.row_slice] = 0
+        else:
+            token_tags[:, segment.row_slice] = 2
+
+    video_denoise_rows = _validate_target_denoise_rows(
+        target_video_denoise_rows,
+        layout.target_video_segment.row_count,
+        "target video denoise rows",
+    )
+    audio_denoise_rows = _validate_target_denoise_rows(
+        target_audio_denoise_rows,
+        layout.target_audio_segment.row_count,
+        "target audio denoise rows",
+    )
+    row_timesteps[:, layout.target_video_segment.row_slice] = torch.where(
+        video_denoise_rows,
+        torch.tensor(video_time, dtype=row_timesteps.dtype),
+        torch.tensor(1.0, dtype=row_timesteps.dtype),
+    )
+    row_timesteps[:, layout.target_audio_segment.row_slice] = torch.where(
+        audio_denoise_rows,
+        torch.tensor(audio_time, dtype=row_timesteps.dtype),
+        torch.tensor(1.0, dtype=row_timesteps.dtype),
+    )
+
+    unique_timesteps, row_timestep_indices = torch.unique(
+        row_timesteps,
+        sorted=True,
+        return_inverse=True,
+    )
+
+    block_adaln_indices = 3 * row_timestep_indices + token_tags
+    block_segments = tuple(_build_modulation_segments(indices) for indices in block_adaln_indices)
+
+    video_timestep_indices = row_timestep_indices[:, layout.target_video_segment.row_slice]
+    audio_timestep_indices = row_timestep_indices[:, layout.target_audio_segment.row_slice]
+
+    return H3TimestepRows(
+        unique_timesteps=unique_timesteps,
+        row_timesteps=row_timesteps,
+        row_timestep_indices=row_timestep_indices,
+        token_tags=token_tags,
+        block_adaln_indices=block_adaln_indices,
+        block_segments=block_segments,
+        video_timestep_index=_timestep_index(
+            unique_timesteps,
+            video_time,
+            "target video",
+            fallback=int(video_timestep_indices[0, 0]),
+        ),
+        audio_timestep_index=_timestep_index(
+            unique_timesteps,
+            audio_time,
+            "target audio",
+            fallback=int(audio_timestep_indices[0, 0]),
+        ),
+        video_timestep_indices=video_timestep_indices,
+        audio_timestep_indices=audio_timestep_indices,
+    )
+
+
+def _unpack_video_rows(rows: torch.Tensor, video: H3VideoGeometry) -> torch.Tensor:
+    unbatched = rows.ndim == 2
+    if unbatched:
+        rows = rows.unsqueeze(0)
+    expected_shape = (video.row_count, VIDEO_CHANNELS * math.prod(VIDEO_PATCH_SIZE))
+    if rows.ndim != 3 or tuple(rows.shape[1:]) != expected_shape:
+        raise ValueError(f"MiniMax-H3 target video rows must be [B,{expected_shape[0]},{expected_shape[1]}]")
+    batch = rows.shape[0]
+    patch_t, patch_h, patch_w = VIDEO_PATCH_SIZE
+    unpacked = rows.reshape(
+        batch,
+        video.frames // patch_t,
+        video.height // patch_h,
+        video.width // patch_w,
+        VIDEO_CHANNELS,
+        patch_t,
+        patch_h,
+        patch_w,
+    )
+    unpacked = unpacked.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(
+        batch,
+        VIDEO_CHANNELS,
+        video.frames,
+        video.height,
+        video.width,
+    )
+    return unpacked[0] if unbatched else unpacked
+
+
+def _unpack_audio_rows(rows: torch.Tensor, audio_frames: int) -> torch.Tensor:
+    unbatched = rows.ndim == 2
+    if unbatched:
+        rows = rows.unsqueeze(0)
+    expected_shape = (audio_frames * STEREO_CHANNELS, AUDIO_CHANNELS)
+    if rows.ndim != 3 or tuple(rows.shape[1:]) != expected_shape:
+        raise ValueError(f"MiniMax-H3 target audio rows must be [B,{expected_shape[0]},{expected_shape[1]}]")
+    batch = rows.shape[0]
+    unpacked = rows.reshape(batch, STEREO_CHANNELS, audio_frames, AUDIO_CHANNELS).permute(0, 3, 1, 2)
+    return unpacked[0] if unbatched else unpacked
+
+
+def unpack_targets(
+    layout: H3PackedLayout,
+    target_video_rows: torch.Tensor,
+    target_audio_rows: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if target_video_rows.ndim != target_audio_rows.ndim or target_video_rows.ndim not in {2, 3}:
+        raise ValueError("MiniMax-H3 target video/audio rows must both be batched or both be unbatched")
+    if target_video_rows.ndim == 3 and target_video_rows.shape[0] != target_audio_rows.shape[0]:
+        raise ValueError("MiniMax-H3 target video/audio rows must have matching batch axes")
+    return (
+        _unpack_video_rows(target_video_rows, layout.target_video),
+        _unpack_audio_rows(target_audio_rows, layout.target_audio_frames),
+    )
diff --git a/musubi_tuner/minimax_h3/sampling.py b/musubi_tuner/minimax_h3/sampling.py
new file mode 100644
index 0000000000000000000000000000000000000000..f61f8a5793d4594d46a9741e97d8d8fa11a17558
--- /dev/null
+++ b/musubi_tuner/minimax_h3/sampling.py
@@ -0,0 +1,396 @@
+# Copyright 2025 The MiniMax authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted for Musubi from Hugging Face Diffusers PR #14355 at commit
+# abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc
+# (schedulers/scheduling_minimax_h3.py and modular_pipelines/minimax_h3/denoise.py).
+# ComfyUI is used only as an independent numerical reference.
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from fractions import Fraction
+from pathlib import Path
+
+import av
+import torch
+
+from musubi_tuner.minimax_h3.packing import H3PackedLayout
+
+
+@dataclass(frozen=True)
+class H3SigmaSchedule:
+    base: torch.Tensor
+    video: torch.Tensor
+    audio: torch.Tensor
+
+
+@dataclass(frozen=True)
+class H3SampleResult:
+    video: torch.Tensor
+    audio: torch.Tensor
+
+
+@dataclass(frozen=True)
+class H3DecodedAV:
+    video: torch.Tensor
+    audio: torch.Tensor
+    fps: int = 24
+    sample_rate: int = 32000
+
+
+def _validate_shift(value: float, label: str) -> float:
+    value = float(value)
+    if not 0.01 <= value <= 100.0:
+        raise ValueError(f"MiniMax-H3 {label} shift must be in [0.01,100.0], got {value}")
+    return value
+
+
+def _shift(base: torch.Tensor, value: float) -> torch.Tensor:
+    return value * base / (1.0 + (value - 1.0) * base)
+
+
+def build_shifted_schedule(
+    steps: int,
+    *,
+    video_shift: float = 12.0,
+    audio_shift: float = 3.0,
+    device: torch.device | str | None = None,
+) -> H3SigmaSchedule:
+    if not isinstance(steps, int) or steps <= 0:
+        raise ValueError(f"MiniMax-H3 sampling steps must be a positive integer, got {steps}")
+    video_shift = _validate_shift(video_shift, "video")
+    audio_shift = _validate_shift(audio_shift, "audio")
+    base = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64, device=device)
+    return H3SigmaSchedule(base=base, video=_shift(base, video_shift), audio=_shift(base, audio_shift))
+
+
+def initialize_target_latents(
+    *,
+    video_shape: Sequence[int],
+    audio_shape: Sequence[int],
+    seed: int,
+    device: torch.device | str,
+    video_dtype: torch.dtype = torch.float16,
+    audio_dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    video_shape = tuple(int(value) for value in video_shape)
+    audio_shape = tuple(int(value) for value in audio_shape)
+    if len(video_shape) != 5 or video_shape[1] != 24:
+        raise ValueError(f"MiniMax-H3 target video noise shape must be [B,24,F,H,W], got {video_shape}")
+    if len(audio_shape) != 4 or audio_shape[1:3] != (32, 2) or audio_shape[0] != video_shape[0]:
+        raise ValueError(f"MiniMax-H3 target audio noise shape must be [B,32,2,A], got {audio_shape}")
+    generator = torch.Generator(device="cpu").manual_seed(int(seed))
+    video = torch.randn(video_shape, generator=generator, dtype=torch.float32, device="cpu").to(device=device, dtype=video_dtype)
+    audio = torch.randn(audio_shape, generator=generator, dtype=torch.float32, device="cpu").to(device=device, dtype=audio_dtype)
+    return video, audio
+
+
+def _augment_condition_group(
+    tensors: Sequence[torch.Tensor],
+    *,
+    seed: int,
+    clean: float,
+    device: torch.device | str,
+) -> tuple[torch.Tensor, ...]:
+    moved = tuple(tensor.to(device) for tensor in tensors)
+    if clean == 1.0:
+        return moved
+    augmented = []
+    for tensor in moved:
+        generator = torch.Generator(device="cpu").manual_seed(int(seed))
+        noise = torch.randn(tuple(tensor.shape), generator=generator, dtype=torch.float32, device="cpu").to(
+            device=tensor.device, dtype=tensor.dtype
+        )
+        augmented.append(clean * tensor + (1.0 - clean) * noise)
+    return tuple(augmented)
+
+
+def augment_condition_latents(
+    visual_conditions: Sequence[torch.Tensor],
+    audio_conditions: Sequence[torch.Tensor],
+    *,
+    seed: int,
+    visual_clean: float = 0.999,
+    audio_clean: float = 1.0,
+    device: torch.device | str,
+) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]:
+    if not 0.0 <= visual_clean <= 1.0:
+        raise ValueError(f"MiniMax-H3 visual condition clean coefficient must be in [0,1], got {visual_clean}")
+    if not 0.0 <= audio_clean <= 1.0:
+        raise ValueError(f"MiniMax-H3 audio condition clean coefficient must be in [0,1], got {audio_clean}")
+    return (
+        _augment_condition_group(visual_conditions, seed=seed, clean=visual_clean, device=device),
+        _augment_condition_group(audio_conditions, seed=seed + 1, clean=audio_clean, device=device),
+    )
+
+
+def _prepare_retake_latents(
+    input_latents: torch.Tensor | None,
+    denoise_mask: torch.Tensor | None,
+    target_latents: torch.Tensor,
+    *,
+    label: str,
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    if (input_latents is None) != (denoise_mask is None):
+        raise ValueError(f"MiniMax-H3 {label} retake requires both input latents and a denoise mask")
+    if input_latents is None:
+        return None, None
+    input_latents = torch.as_tensor(input_latents)
+    if input_latents.shape != target_latents.shape:
+        raise ValueError(
+            f"MiniMax-H3 {label} retake input shape {tuple(input_latents.shape)} "
+            f"does not match target {tuple(target_latents.shape)}"
+        )
+    input_latents = input_latents.to(device=target_latents.device, dtype=target_latents.dtype)
+    denoise_mask = torch.as_tensor(denoise_mask, device=target_latents.device)
+    if denoise_mask.is_complex():
+        raise ValueError(f"MiniMax-H3 {label} denoise mask must be real")
+    if denoise_mask.is_floating_point() and not torch.isfinite(denoise_mask).all():
+        raise ValueError(f"MiniMax-H3 {label} denoise mask must be finite")
+    if not torch.all((denoise_mask == 0) | (denoise_mask == 1)):
+        raise ValueError(f"MiniMax-H3 {label} denoise mask must contain only 0 (fixed) and 1 (regenerate)")
+    try:
+        denoise_mask = torch.broadcast_to(denoise_mask, target_latents.shape)
+    except RuntimeError as error:
+        raise ValueError(
+            f"MiniMax-H3 {label} denoise mask shape {tuple(denoise_mask.shape)} "
+            f"is not broadcastable to {tuple(target_latents.shape)}"
+        ) from error
+    return input_latents, denoise_mask.to(torch.bool)
+
+
+@torch.no_grad()
+def sample_joint_av(
+    transformer,
+    *,
+    layout: H3PackedLayout,
+    text_hidden_states: torch.Tensor,
+    text_token_tags: torch.Tensor,
+    initial_video: torch.Tensor,
+    initial_audio: torch.Tensor,
+    steps: int,
+    video_shift: float = 12.0,
+    audio_shift: float = 3.0,
+    visual_condition_latents: Sequence[torch.Tensor] = (),
+    audio_condition_latents: Sequence[torch.Tensor] = (),
+    visual_condition_clean: float = 0.999,
+    audio_condition_clean: float = 1.0,
+    input_video_latents: torch.Tensor | None = None,
+    video_denoise_mask: torch.Tensor | None = None,
+    input_audio_latents: torch.Tensor | None = None,
+    audio_denoise_mask: torch.Tensor | None = None,
+    step_callback: Callable[[int, int], None] | None = None,
+) -> H3SampleResult:
+    if initial_video.ndim != 5 or tuple(initial_video.shape[2:]) != (
+        layout.target_video.frames,
+        layout.target_video.height,
+        layout.target_video.width,
+    ):
+        raise ValueError("MiniMax-H3 initial video noise does not match the packed layout")
+    if initial_audio.ndim != 4 or tuple(initial_audio.shape[1:]) != (32, 2, layout.target_audio_frames):
+        raise ValueError("MiniMax-H3 initial audio noise does not match the packed layout")
+    if initial_video.shape[0] != initial_audio.shape[0]:
+        raise ValueError("MiniMax-H3 initial video and audio batch sizes differ")
+    if initial_video.shape[0] != 1:
+        raise ValueError(f"MiniMax-H3 R1 requires batch_size=1, got {initial_video.shape[0]}")
+    if text_hidden_states.shape[:2] != (initial_video.shape[0], layout.text_length):
+        raise ValueError("MiniMax-H3 text hidden states do not match the sampling layout")
+    if text_token_tags.shape != (initial_video.shape[0], layout.text_length):
+        raise ValueError("MiniMax-H3 text token tags must preserve the [B,L] axes")
+
+    input_video_latents, video_denoise_mask = _prepare_retake_latents(
+        input_video_latents,
+        video_denoise_mask,
+        initial_video,
+        label="video",
+    )
+    input_audio_latents, audio_denoise_mask = _prepare_retake_latents(
+        input_audio_latents,
+        audio_denoise_mask,
+        initial_audio,
+        label="audio",
+    )
+
+    schedule = build_shifted_schedule(
+        steps,
+        video_shift=video_shift,
+        audio_shift=audio_shift,
+        device=initial_video.device,
+    )
+    video = initial_video
+    audio = initial_audio
+    for index in range(steps):
+        sigma_video = schedule.video[index].to(dtype=torch.float32)
+        sigma_audio = schedule.audio[index].to(dtype=torch.float32)
+        model_video = (
+            video
+            if video_denoise_mask is None
+            else torch.where(video_denoise_mask, video, input_video_latents)
+        )
+        model_audio = (
+            audio
+            if audio_denoise_mask is None
+            else torch.where(audio_denoise_mask, audio, input_audio_latents)
+        )
+        retake_kwargs = {}
+        if video_denoise_mask is not None:
+            retake_kwargs["video_denoise_mask"] = video_denoise_mask
+        if audio_denoise_mask is not None:
+            retake_kwargs["audio_denoise_mask"] = audio_denoise_mask
+        prediction = transformer(
+            video_latents=model_video,
+            audio_latents=model_audio,
+            text_hidden_states=text_hidden_states,
+            text_token_tags=text_token_tags,
+            layout=layout,
+            model_t_video=1.0 - sigma_video,
+            model_t_audio=1.0 - sigma_audio,
+            visual_condition_latents=visual_condition_latents,
+            audio_condition_latents=audio_condition_latents,
+            visual_condition_clean=visual_condition_clean,
+            audio_condition_clean=audio_condition_clean,
+            **retake_kwargs,
+        )
+        if prediction.video.shape != video.shape or prediction.audio.shape != audio.shape:
+            raise ValueError("MiniMax-H3 transformer predictions do not match the target latent shapes")
+        prediction_video = prediction.video
+        prediction_audio = prediction.audio
+        if video_denoise_mask is not None:
+            fixed_prediction = (input_video_latents - video) / sigma_video.to(video)
+            prediction_video = torch.where(video_denoise_mask, prediction_video, fixed_prediction)
+        if audio_denoise_mask is not None:
+            fixed_prediction = (input_audio_latents - audio) / sigma_audio.to(audio)
+            prediction_audio = torch.where(audio_denoise_mask, prediction_audio, fixed_prediction)
+        video_delta = (schedule.video[index] - schedule.video[index + 1]).to(video)
+        audio_delta = (schedule.audio[index] - schedule.audio[index + 1]).to(audio)
+        video = video + video_delta * prediction_video
+        audio = audio + audio_delta * prediction_audio
+        if step_callback is not None:
+            step_callback(index + 1, steps)
+    return H3SampleResult(video=video, audio=audio)
+
+
+@torch.no_grad()
+def decode_joint_av(
+    video_vae,
+    audio_vae,
+    sample: H3SampleResult,
+    *,
+    frame_count: int,
+    fps: int = 24,
+    sample_rate: int = 32000,
+) -> H3DecodedAV:
+    decoded_video = video_vae.decode(sample.video)
+    decoded_audio = audio_vae.decode(sample.audio)
+    return synchronize_decoded_av(
+        decoded_video,
+        decoded_audio,
+        frame_count=frame_count,
+        fps=fps,
+        sample_rate=sample_rate,
+    )
+
+
+def synchronize_decoded_av(
+    decoded_video: torch.Tensor,
+    decoded_audio: torch.Tensor,
+    *,
+    frame_count: int,
+    fps: int = 24,
+    sample_rate: int = 32000,
+) -> H3DecodedAV:
+    if frame_count <= 0 or fps <= 0 or sample_rate <= 0:
+        raise ValueError("MiniMax-H3 decode frame count, fps, and sample rate must be positive")
+    if decoded_video.ndim != 5 or decoded_video.shape[0] != 1 or decoded_video.shape[1] != 3:
+        raise ValueError(f"MiniMax-H3 video VAE must decode [1,3,F,H,W], got {tuple(decoded_video.shape)}")
+    if decoded_audio.ndim != 3 or decoded_audio.shape[:2] != (1, 2):
+        raise ValueError(f"MiniMax-H3 audio VAE must decode [1,2,L], got {tuple(decoded_audio.shape)}")
+
+    planned_audio_samples = round(Fraction(frame_count * sample_rate, fps))
+    video_audio_samples = round(Fraction(decoded_video.shape[2] * sample_rate, fps))
+    common_audio_samples = min(decoded_audio.shape[-1], planned_audio_samples, video_audio_samples)
+    if common_audio_samples <= 0:
+        raise ValueError("MiniMax-H3 decoded audio/video duration is empty")
+    common_video_frames = min(
+        frame_count,
+        decoded_video.shape[2],
+        max(1, round(Fraction(common_audio_samples * fps, sample_rate))),
+    )
+
+    decoded_video = decoded_video[0, :, :common_video_frames].detach().cpu()
+    video_chunks = []
+    for chunk in decoded_video.split(16, dim=1):
+        chunk = chunk.float().clamp(-1.0, 1.0)
+        video_chunks.append(((chunk + 1.0) * 127.5).round().to(torch.uint8).permute(1, 2, 3, 0))
+    video = torch.cat(video_chunks).contiguous()
+    audio = decoded_audio[0, :, :common_audio_samples].detach().cpu().float().clamp(-1.0, 1.0).contiguous()
+    return H3DecodedAV(video=video, audio=audio, fps=fps, sample_rate=sample_rate)
+
+
+def mux_audio_video(
+    video: torch.Tensor,
+    audio: torch.Tensor,
+    output_path: str | Path,
+    *,
+    fps: int,
+    sample_rate: int,
+) -> None:
+    if video.ndim != 4 or video.shape[-1] != 3 or video.dtype != torch.uint8:
+        raise ValueError(f"MiniMax-H3 mux video must be uint8 [F,H,W,3], got {tuple(video.shape)} {video.dtype}")
+    if audio.ndim != 2 or audio.shape[0] != 2:
+        raise ValueError(f"MiniMax-H3 mux audio must be stereo [2,L], got {tuple(audio.shape)}")
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with av.open(str(output_path), mode="w") as container:
+        video_stream = container.add_stream("libx264", rate=fps)
+        video_stream.width = video.shape[2]
+        video_stream.height = video.shape[1]
+        video_stream.pix_fmt = "yuv420p"
+        audio_stream = container.add_stream("aac", rate=sample_rate)
+        audio_stream.layout = "stereo"
+
+        for pixels in video:
+            frame = av.VideoFrame.from_ndarray(pixels.numpy(), format="rgb24")
+            for packet in video_stream.encode(frame):
+                container.mux(packet)
+        for packet in video_stream.encode():
+            container.mux(packet)
+
+        for start in range(0, audio.shape[1], 1024):
+            samples = audio[:, start : start + 1024].numpy()
+            frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout="stereo")
+            frame.sample_rate = sample_rate
+            frame.pts = start
+            frame.time_base = Fraction(1, sample_rate)
+            for packet in audio_stream.encode(frame):
+                container.mux(packet)
+        for packet in audio_stream.encode():
+            container.mux(packet)
+
+
+def write_joint_av(
+    decoded: H3DecodedAV,
+    output_path: str | Path,
+    *,
+    muxer: Callable[..., None] = mux_audio_video,
+) -> None:
+    muxer(
+        decoded.video,
+        decoded.audio,
+        Path(output_path),
+        fps=decoded.fps,
+        sample_rate=decoded.sample_rate,
+    )
diff --git a/musubi_tuner/minimax_h3/text_encoder.py b/musubi_tuner/minimax_h3/text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..8547551a5a8ed335e2e00adc4b57e12a10d987ed
--- /dev/null
+++ b/musubi_tuner/minimax_h3/text_encoder.py
@@ -0,0 +1,384 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted for Musubi from Hugging Face Diffusers PR #14355 at commit
+# abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc
+# (modular_pipelines/minimax_h3/encoders.py and packing_ref2va.py).
+# ComfyUI is used only as an independent numerical reference.
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+import hashlib
+import json
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import torch
+
+from musubi_tuner.minimax_h3.checkpoint import load_safetensors_module, resolve_safetensors_files
+from musubi_tuner.minimax_h3.media import H3Record, H3Task
+
+
+TEXT_WIDTH = 5120
+MAX_TEXT_ROWS = 32768
+LAYER_50_HIDDEN_STATE_INDEX = 50
+VISION_START_TOKEN_ID = 151652
+VISION_END_TOKEN_ID = 151653
+IMAGE_TOKEN_ID = 151655
+VIDEO_TOKEN_ID = 151656
+IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
+VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
+DEFAULT_PROCESSOR_ID = "Qwen/Qwen3-VL-32B-Instruct"
+
+
+@dataclass(frozen=True)
+class H3TextVisual:
+    frames: torch.Tensor
+    timestamps: tuple[float, ...] | None = None
+
+    def __post_init__(self) -> None:
+        if self.frames.ndim != 4 or self.frames.shape[-1] < 3 or self.frames.shape[0] == 0:
+            raise ValueError(f"MiniMax-H3 text visual must be [T,H,W,C], got {tuple(self.frames.shape)}")
+        if self.timestamps is not None and len(self.timestamps) != self.frames.shape[0]:
+            raise ValueError("MiniMax-H3 text visual timestamps must match its frame count")
+
+
+@dataclass(frozen=True)
+class H3Presentation:
+    text: str
+    images: tuple[torch.Tensor, ...] = ()
+    videos: tuple[torch.Tensor, ...] = ()
+    processor_text: str | None = None
+
+
+def _require_visual(visuals: Mapping[object, H3TextVisual], key: object, label: str) -> H3TextVisual:
+    try:
+        return visuals[key]
+    except KeyError as error:
+        raise ValueError(f"MiniMax-H3 presentation is missing {label} visual data") from error
+
+
+def build_presentation(
+    record: H3Record,
+    task: H3Task,
+    visuals: Mapping[object, H3TextVisual] | None = None,
+) -> H3Presentation:
+    if task not in {"t2va", "fl2va", "ref2va"}:
+        raise ValueError(f"Unsupported MiniMax-H3 task: {task}")
+    visuals = visuals or {}
+    if task != "ref2va" and record.references:
+        raise ValueError(f"MiniMax-H3 task {task} does not accept references")
+    if task == "t2va":
+        return H3Presentation(text=record.caption, processor_text=record.caption)
+
+    parts = []
+    processor_parts = []
+    images = []
+    videos = []
+    if task == "fl2va":
+        for index, key in enumerate(("first", "last"), start=1):
+            visual = _require_visual(visuals, key, f"FL2VA {key}")
+            if visual.frames.shape[0] != 1:
+                raise ValueError(f"MiniMax-H3 FL2VA {key} visual must contain exactly one frame")
+            part = f"<Picture {index}>: {IMAGE_PLACEHOLDER}"
+            parts.append(part)
+            processor_parts.append(part)
+            images.append(visual.frames[0])
+        parts.append(record.caption)
+        processor_parts.append(record.caption)
+        return H3Presentation(
+            text="".join(parts),
+            images=tuple(images),
+            processor_text="".join(processor_parts),
+        )
+
+    counters = {"image": 0, "audio": 0, "video": 0}
+    for reference in record.references:
+        if reference.type == "image":
+            counters["image"] += 1
+            visual = _require_visual(visuals, reference.path, f"reference image {reference.path}")
+            if visual.frames.shape[0] != 1:
+                raise ValueError(f"MiniMax-H3 reference image must contain one frame: {reference.path}")
+            part = f"<Picture {counters['image']}>: {IMAGE_PLACEHOLDER}"
+            parts.append(part)
+            processor_parts.append(part)
+            images.append(visual.frames[0])
+            continue
+
+        if reference.type == "audio":
+            counters["audio"] += 1
+            part = f"<Audio {counters['audio']}>: "
+            parts.append(part)
+            processor_parts.append(part)
+            continue
+
+        if reference.type != "video":
+            raise ValueError(f"Unsupported MiniMax-H3 reference type: {reference.type}")
+        if reference.audio is not None:
+            counters["audio"] += 1
+            part = f"<Audio {counters['audio']}>: "
+            parts.append(part)
+            processor_parts.append(part)
+        counters["video"] += 1
+        part = f"<Video {counters['video']}>: "
+        parts.append(part)
+        processor_parts.append(part)
+        processor_parts.append(VIDEO_PLACEHOLDER)
+
+        visual = _require_visual(visuals, reference.path, f"reference video {reference.path}")
+        frames = visual.frames
+        timestamps = list(visual.timestamps) if visual.timestamps is not None else [index / 2.0 for index in range(len(frames))]
+        if len(frames) % 2:
+            frames = torch.cat((frames, frames[-1:]), dim=0)
+            timestamps.append(timestamps[-1])
+        for index in range(0, len(frames), 2):
+            block_timestamp = (timestamps[index] + timestamps[index + 1]) / 2.0
+            parts.append(f"<{block_timestamp:.1f} seconds>{VIDEO_PLACEHOLDER}")
+        videos.append(visual.frames)
+
+    parts.append(record.caption)
+    processor_parts.append(record.caption)
+    return H3Presentation(
+        text="".join(parts),
+        images=tuple(images),
+        videos=tuple(videos),
+        processor_text="".join(processor_parts),
+    )
+
+
+def build_token_tags(processed: Mapping[str, Any]) -> torch.Tensor:
+    if "input_ids" not in processed:
+        raise ValueError("MiniMax-H3 processor output has no input_ids")
+    input_ids = torch.as_tensor(processed["input_ids"])
+    if input_ids.ndim == 2:
+        if input_ids.shape[0] != 1:
+            raise ValueError("MiniMax-H3 text caching processes one presentation at a time")
+        input_ids = input_ids[0]
+    if input_ids.ndim != 1:
+        raise ValueError(f"MiniMax-H3 input_ids must be [L] or [1,L], got {tuple(input_ids.shape)}")
+
+    tags = torch.ones(input_ids.shape[0], dtype=torch.int64, device=input_ids.device)
+    open_start = None
+    for index, token_id in enumerate(input_ids.tolist()):
+        if token_id == VISION_START_TOKEN_ID:
+            if open_start is not None:
+                raise ValueError("MiniMax-H3 processor produced nested vision-start tokens")
+            open_start = index
+        elif token_id == VISION_END_TOKEN_ID:
+            if open_start is None:
+                raise ValueError("MiniMax-H3 processor produced an unmatched vision-end token")
+            interior = input_ids[open_start + 1 : index]
+            if not torch.any((interior == IMAGE_TOKEN_ID) | (interior == VIDEO_TOKEN_ID)):
+                raise ValueError("MiniMax-H3 vision span contains no expanded image or video rows")
+            tags[open_start : index + 1] = 0
+            open_start = None
+    if open_start is not None:
+        raise ValueError("MiniMax-H3 processor produced an unmatched vision-start token")
+    return tags.cpu()
+
+
+def _text_size_error(row_count: int, token_tags: torch.Tensor) -> ValueError:
+    vision_rows = int((token_tags == 0).sum().item())
+    text_rows = row_count - vision_rows
+    payload_mib = row_count * TEXT_WIDTH * 2 / (1024**2)
+    return ValueError(
+        f"MiniMax-H3 text presentation has {row_count} rows, exceeds {MAX_TEXT_ROWS}; "
+        f"vision_rows={vision_rows}, text_rows={text_rows}, estimated BF16 payload={payload_mib:.1f} MiB"
+    )
+
+
+def validate_text_rows(hidden_states: torch.Tensor, token_tags: torch.Tensor) -> None:
+    if hidden_states.ndim != 2 or hidden_states.shape[1] != TEXT_WIDTH:
+        raise ValueError(f"Expected MiniMax-H3 hidden states [L,{TEXT_WIDTH}], got {tuple(hidden_states.shape)}")
+    if token_tags.dtype != torch.int64 or token_tags.shape != (hidden_states.shape[0],):
+        raise ValueError("MiniMax-H3 token tags must be int64 [L]")
+    if hidden_states.shape[0] > MAX_TEXT_ROWS:
+        raise _text_size_error(hidden_states.shape[0], token_tags)
+    if not torch.all((token_tags == 0) | (token_tags == 1)):
+        raise ValueError("MiniMax-H3 token tags may contain only 0 and 1")
+
+
+def _language_layers(model) -> list:
+    language_model = getattr(model, "language_model", None)
+    layers = getattr(language_model, "layers", None)
+    if layers is None:
+        raise ValueError("MiniMax-H3 Qwen3-VL model has no language_model.layers")
+    return layers
+
+
+def extract_layer_50_pre_norm(output, model) -> torch.Tensor:
+    layers = _language_layers(model)
+    if len(layers) < LAYER_50_HIDDEN_STATE_INDEX:
+        raise ValueError(f"MiniMax-H3 Qwen3-VL needs at least 50 layers, got {len(layers)}")
+
+    captured = getattr(output, "h3_layer_50_pre_norm", None)
+    if captured is None:
+        captured = getattr(model, "_h3_layer_50_pre_norm", None)
+    if captured is not None:
+        return captured
+
+    hidden_states = getattr(output, "hidden_states", None)
+    if hidden_states is None or len(hidden_states) <= LAYER_50_HIDDEN_STATE_INDEX:
+        raise ValueError("MiniMax-H3 layer-50 pre-norm state was not captured")
+    return hidden_states[LAYER_50_HIDDEN_STATE_INDEX]
+
+
+def normalize_h3_text_encoder_key(key: str) -> str:
+    if key.startswith("model."):
+        return "language_model." + key.removeprefix("model.")
+    return key
+
+
+def load_h3_processor(path: str = DEFAULT_PROCESSOR_ID, *, revision: str | None = None):
+    from transformers import AutoProcessor
+
+    return AutoProcessor.from_pretrained(path, revision=revision)
+
+
+def load_h3_text_encoder(
+    checkpoint_path: str | Path,
+    *,
+    processor_path: str = DEFAULT_PROCESSOR_ID,
+    revision: str | None = None,
+    device: str | torch.device,
+    dtype: torch.dtype = torch.bfloat16,
+    disable_mmap: bool = False,
+):
+    from transformers import Qwen3VLConfig, Qwen3VLModel
+
+    config = Qwen3VLConfig.from_pretrained(processor_path, revision=revision)
+    if config.text_config.hidden_size != TEXT_WIDTH:
+        raise ValueError(f"MiniMax-H3 Qwen3-VL hidden size must be {TEXT_WIDTH}, got {config.text_config.hidden_size}")
+    config.text_config.num_hidden_layers = LAYER_50_HIDDEN_STATE_INDEX
+    config.text_config.use_cache = False
+
+    def factory():
+        model = Qwen3VLModel(config)
+        del model.language_model.norm
+        return model
+
+    return load_safetensors_module(
+        factory,
+        resolve_safetensors_files(checkpoint_path),
+        device=device,
+        dtype=dtype,
+        key_transform=normalize_h3_text_encoder_key,
+        disable_mmap=disable_mmap,
+    )
+
+
+class _Layer50Captured(Exception):
+    pass
+
+
+def _processor_value(value):
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu().numpy()
+    return np.asarray(value)
+
+
+@torch.no_grad()
+def encode_h3_presentation(processor, model, presentation: H3Presentation) -> tuple[torch.Tensor, torch.Tensor]:
+    processor_args = {
+        "text": [presentation.processor_text or presentation.text],
+        "return_tensors": "pt",
+        "padding": False,
+    }
+    if presentation.images:
+        processor_args["images"] = [_processor_value(image) for image in presentation.images]
+    if presentation.videos:
+        processor_args["videos"] = [_processor_value(video) for video in presentation.videos]
+        processor_args["video_metadata"] = [
+            {
+                "total_num_frames": int(video.shape[0]),
+                "fps": 2.0,
+                "duration": float(video.shape[0]) / 2.0,
+                "frames_indices": list(range(video.shape[0])),
+                "height": int(video.shape[1]),
+                "width": int(video.shape[2]),
+            }
+            for video in presentation.videos
+        ]
+        processor_args["do_sample_frames"] = False
+    processed = processor(**processor_args)
+    token_tags = build_token_tags(processed)
+    row_count = token_tags.shape[0]
+    if row_count > MAX_TEXT_ROWS:
+        raise _text_size_error(row_count, token_tags)
+
+    layers = _language_layers(model)
+    if len(layers) != LAYER_50_HIDDEN_STATE_INDEX:
+        raise ValueError(f"Released MiniMax-H3 text encoder must contain exactly 50 layers, got {len(layers)}")
+    captured = {}
+
+    def capture_layer_50(module, args, output):
+        del module, args
+        captured["hidden"] = output[0] if isinstance(output, tuple) else output
+        raise _Layer50Captured
+
+    handle = layers[LAYER_50_HIDDEN_STATE_INDEX - 1].register_forward_hook(capture_layer_50)
+    device = next(model.parameters()).device
+    model_inputs = {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in processed.items()
+        if key != "token_type_ids"
+    }
+    try:
+        model(**model_inputs, use_cache=False)
+    except _Layer50Captured:
+        pass
+    finally:
+        handle.remove()
+    if "hidden" not in captured:
+        raise RuntimeError("MiniMax-H3 text encoder did not reach layer 50")
+
+    output = SimpleNamespace(h3_layer_50_pre_norm=captured["hidden"], hidden_states=None, last_hidden_state=None)
+    hidden_states = extract_layer_50_pre_norm(output, model)
+    if hidden_states.ndim != 3 or hidden_states.shape[0] != 1:
+        raise ValueError(f"MiniMax-H3 layer-50 output must be [1,L,{TEXT_WIDTH}], got {tuple(hidden_states.shape)}")
+    hidden_states = hidden_states[0]
+    validate_text_rows(hidden_states, token_tags)
+    return hidden_states, token_tags
+
+
+def processor_fingerprint(processor) -> str:
+    tokenizer = getattr(processor, "tokenizer", None)
+    payload = {
+        "processor_class": type(processor).__name__,
+        "processor_name": getattr(processor, "name_or_path", None),
+        "tokenizer_class": type(tokenizer).__name__ if tokenizer is not None else None,
+        "tokenizer_name": getattr(tokenizer, "name_or_path", None),
+        "tokenizer_size": len(tokenizer) if tokenizer is not None else None,
+        "commit": getattr(processor, "_commit_hash", None) or (getattr(tokenizer, "init_kwargs", {}) or {}).get("_commit_hash"),
+        "vision_tokens": [VISION_START_TOKEN_ID, VISION_END_TOKEN_ID, IMAGE_TOKEN_ID, VIDEO_TOKEN_ID],
+    }
+    encoded = json.dumps(payload, sort_keys=True, ensure_ascii=True, separators=(",", ":")).encode("utf-8")
+    return f"sha256:{hashlib.sha256(encoded).hexdigest()}"
+
+
+def presentation_fingerprint(presentation: H3Presentation, media_fingerprints: Mapping[Path, str]) -> str:
+    payload = {
+        "text": presentation.text,
+        "processor_text": presentation.processor_text,
+        "image_shapes": [list(image.shape) for image in presentation.images],
+        "video_shapes": [list(video.shape) for video in presentation.videos],
+        "media": dict(sorted((str(Path(path).resolve()), value) for path, value in media_fingerprints.items())),
+        "format": "minimax-h3-non-chat-v1",
+    }
+    encoded = json.dumps(payload, sort_keys=True, ensure_ascii=True, separators=(",", ":")).encode("utf-8")
+    return f"sha256:{hashlib.sha256(encoded).hexdigest()}"
diff --git a/musubi_tuner/minimax_h3/video_vae.py b/musubi_tuner/minimax_h3/video_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..67f6d12e0e1d1b5f5897e8d0e8a739a8346a553c
--- /dev/null
+++ b/musubi_tuner/minimax_h3/video_vae.py
@@ -0,0 +1,686 @@
+# Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Adapted for Musubi from Hugging Face Diffusers PR #14355 at commit
+# abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc. Musubi keeps the published
+# checkpoint names and adds cache-specific posterior sampling wrappers.
+# ComfyUI is used only as an independent numerical reference.
+
+import hashlib
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+LATENTS_MEAN = [
+    0.858090341091156,
+    -0.9606591463088989,
+    1.0661640167236328,
+    -0.5090325474739075,
+    -0.2727581858634949,
+    -1.3675414323806763,
+    -0.2553254961967468,
+    -0.26907554268836975,
+    -0.5376840829849243,
+    -0.0464097298681736,
+    0.6657370328903198,
+    0.19690127670764923,
+    -0.5460608005523682,
+    -0.4035342037677765,
+    -0.23683024942874908,
+    0.25928452610969543,
+    -0.30133944749832153,
+    0.211341992020607,
+    -1.1206848621368408,
+    0.3581933379173279,
+    -0.04225143790245056,
+    0.2604829967021942,
+    0.22864092886447906,
+    0.7056031823158264,
+]
+
+LATENTS_STD = [
+    1.2223774194717407,
+    1.2767263650894165,
+    1.68317747116088865,
+    1.7549455165863037,
+    1.5636216402053833,
+    2.194143533706665,
+    0.96531379222869875,
+    1.05698859691619875,
+    0.841948926448822,
+    0.7729952931404114,
+    1.8955937623977661,
+    0.946841835975647,
+    0.7996809482574463,
+    0.44988900423049925,
+    0.7197399735450745,
+    0.69362932443618775,
+    2.961095094680786,
+    2.7694199085235595,
+    3.0496184825897215,
+    2.1088054180145265,
+    3.276226282119751,
+    3.1627357006073,
+    2.28168129920959475,
+    2.6127843856811525,
+]
+
+
+class CausalConv3d(nn.Conv3d):
+    """Spatially reflected, temporally causal 3D convolution."""
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
+        super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=0)
+        self.causal_padding = (padding,) * 3 if isinstance(padding, int) else tuple(padding)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        temporal, height, width = self.causal_padding
+        if height or width:
+            hidden_states = F.pad(hidden_states, (width, width, height, height, 0, 0), mode="reflect")
+        if temporal:
+            hidden_states = F.pad(hidden_states, (0, 0, 0, 0, temporal * 2, 0), mode="constant")
+        return F.conv3d(
+            hidden_states,
+            self.weight,
+            self.bias,
+            stride=self.stride,
+            padding=0,
+            dilation=self.dilation,
+            groups=self.groups,
+        )
+
+
+class TemporalIsolatedGroupNorm(nn.GroupNorm):
+    """Compute group-normalization statistics independently for each frame."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if hidden_states.ndim != 5:
+            return super().forward(hidden_states)
+        batch_size, channels, frames, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 1, 3, 4).contiguous()
+        hidden_states = hidden_states.view(batch_size * frames, channels, 1, height, width)
+        hidden_states = super().forward(hidden_states)
+        hidden_states = hidden_states.view(batch_size, frames, channels, height, width)
+        return hidden_states.permute(0, 2, 1, 3, 4).contiguous()
+
+
+def group_norm_3d(num_channels: int) -> TemporalIsolatedGroupNorm:
+    return TemporalIsolatedGroupNorm(32, num_channels, eps=1e-6, affine=True)
+
+
+class Downsample3D(nn.Module):
+    def __init__(self, in_channels, out_channels, time_stride=1, space_stride=2):
+        super().__init__()
+        self.space_stride = space_stride
+        self.conv = CausalConv3d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=(time_stride, space_stride, space_stride),
+            padding=(1, 0, 0),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.space_stride == 2:
+            hidden_states = F.pad(hidden_states, (0, 1, 0, 1, 0, 0), mode="reflect")
+        return self.conv(hidden_states)
+
+
+class ResnetBlock3D(nn.Module):
+    def __init__(self, in_channels, out_channels=None):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.norm1 = group_norm_3d(in_channels)
+        self.conv1 = CausalConv3d(in_channels, self.out_channels, kernel_size=3, padding=1)
+        self.norm2 = group_norm_3d(self.out_channels)
+        self.conv2 = CausalConv3d(self.out_channels, self.out_channels, kernel_size=3, padding=1)
+        self.nin_shortcut = (
+            CausalConv3d(in_channels, self.out_channels, kernel_size=1) if in_channels != self.out_channels else None
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.conv1(F.silu(self.norm1(hidden_states)))
+        hidden_states = self.conv2(F.silu(self.norm2(hidden_states)))
+        if self.nin_shortcut is not None:
+            residual = self.nin_shortcut(residual)
+        return residual + hidden_states
+
+
+class EncoderFCN3D(nn.Module):
+    def __init__(self, ch, ch_mult, space_down, time_down, num_res_blocks, in_channels, z_channels, double_z=True):
+        super().__init__()
+        num_levels = len(ch_mult)
+        layers_per_level = [num_res_blocks] * num_levels if isinstance(num_res_blocks, int) else list(num_res_blocks)
+        block_channels = [ch * multiplier for multiplier in ch_mult]
+        input_channels = [block_channels[0], *block_channels[:-1]]
+
+        self.num_levels = num_levels
+        self.num_res_blocks = layers_per_level
+        self.conv_in = CausalConv3d(in_channels, input_channels[0], kernel_size=3, padding=1)
+        self.down = nn.ModuleList()
+        for level in range(num_levels):
+            down = nn.Module()
+            down.block = nn.ModuleList(
+                [
+                    ResnetBlock3D(
+                        input_channels[level] if layer == 0 else block_channels[level],
+                        block_channels[level],
+                    )
+                    for layer in range(layers_per_level[level])
+                ]
+            )
+            if space_down[level] * time_down[level] > 1:
+                down.downsample = Downsample3D(
+                    block_channels[level],
+                    block_channels[level],
+                    time_stride=time_down[level],
+                    space_stride=space_down[level],
+                )
+            self.down.append(down)
+
+        self.norm_out = group_norm_3d(block_channels[-1])
+        output_channels = 2 * z_channels if double_z else z_channels
+        self.conv_out = CausalConv3d(block_channels[-1], output_channels, kernel_size=3, padding=1)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+        for down in self.down:
+            for block in down.block:
+                hidden_states = block(hidden_states)
+            if hasattr(down, "downsample"):
+                hidden_states = down.downsample(hidden_states)
+        hidden_states = F.silu(self.norm_out(hidden_states))
+        return self.conv_out(hidden_states)
+
+
+def create_token_ids(patch_dims, device, dtype=torch.float32):
+    grids = [2.0 * (torch.arange(0.5, size, dtype=dtype, device=device) / size) - 1.0 for size in patch_dims]
+    return torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=-1).flatten(0, len(patch_dims) - 1).unsqueeze(0)
+
+
+class RotaryEmbeddingND(nn.Module):
+    def __init__(self, dim, rotary_base=100.0, n_dim=3):
+        super().__init__()
+        if dim % (2 * n_dim):
+            raise ValueError(f"Rotary dimension {dim} must be divisible by {2 * n_dim}")
+        inv_freq = 1.0 / rotary_base ** torch.arange(0, 1, 2 * n_dim / dim, dtype=torch.float32)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        angles = 2.0 * math.pi * position_ids[:, :, :, None].float() * self.inv_freq[None, None, None, :]
+        angles = angles.flatten(2, 3).tile(2).unsqueeze(2)
+        return angles.cos(), angles.sin()
+
+
+def _apply_rotary_emb(hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    cosine, sine = (value.to(hidden_states.dtype) for value in rotary_emb)
+    rotary_dim = cosine.shape[-1]
+    rotary, passthrough = hidden_states[..., :rotary_dim], hidden_states[..., rotary_dim:]
+    first, second = rotary.chunk(2, dim=-1)
+    rotated = torch.cat((-second, first), dim=-1)
+    return torch.cat((rotary * cosine + rotated * sine, passthrough), dim=-1)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4, bias=True):
+        super().__init__()
+        inner_dim = dim * mult
+        self.w1 = nn.Linear(dim, inner_dim * 2, bias=bias)
+        self.w2 = nn.Linear(inner_dim, dim, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate, value = self.w1(hidden_states).chunk(2, dim=-1)
+        return self.w2(F.silu(gate) * value)
+
+
+class Attention(nn.Module):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5):
+        super().__init__()
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = heads * dim_head
+        self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
+        self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
+        self.to_qkv = nn.Linear(inner_dim, inner_dim * 3, bias=bias)
+        self.to_out = nn.Linear(inner_dim, inner_dim, bias=bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        batch_size, sequence_length, _ = hidden_states.shape
+        qkv = self.to_qkv(hidden_states).view(batch_size, sequence_length, self.heads, 3 * self.dim_head)
+        query, key, value = qkv.chunk(3, dim=-1)
+        query = self.norm_q(query.float()).to(query.dtype)
+        key = self.norm_k(key.float()).to(key.dtype)
+        if rotary_emb is not None:
+            query = _apply_rotary_emb(query, rotary_emb)
+            key = _apply_rotary_emb(key, rotary_emb)
+        hidden_states = F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, sequence_length, -1)
+        return self.to_out(hidden_states)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5):
+        super().__init__()
+        dim = heads * dim_head
+        self.norm1 = nn.RMSNorm(dim, elementwise_affine=True, eps=eps)
+        self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
+        self.scale1 = nn.Parameter(torch.empty(dim))
+        self.norm2 = nn.RMSNorm(dim, elementwise_affine=True, eps=eps)
+        self.ff = FeedForward(dim=dim, bias=bias)
+        self.scale2 = nn.Parameter(torch.empty(dim))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        normalized = F.rms_norm(
+            hidden_states.float(),
+            (hidden_states.shape[-1],),
+            self.norm1.weight.float(),
+            self.norm1.eps,
+        ).to(hidden_states.dtype)
+        hidden_states = hidden_states + self.attn(normalized, rotary_emb) * self.scale1
+        normalized = F.rms_norm(
+            hidden_states.float(),
+            (hidden_states.shape[-1],),
+            self.norm2.weight.float(),
+            self.norm2.eps,
+        ).to(hidden_states.dtype)
+        return hidden_states + self.ff(normalized) * self.scale2
+
+
+class ViT3DDecoder(nn.Module):
+    def __init__(
+        self,
+        patch_size=16,
+        patch_size_t=4,
+        in_channels=24,
+        out_channels=3,
+        num_layers=36,
+        heads=32,
+        dim_head=64,
+        rope_theta=100.0,
+        rope_dim_ratio=0.75,
+        bias=True,
+        eps=1e-5,
+        num_register_tokens=4,
+    ):
+        super().__init__()
+        dim = heads * dim_head
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.out_channels = out_channels
+        self.num_register_tokens = num_register_tokens
+        self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, n_dim=3)
+        self.x_embedder = nn.Linear(in_channels, dim)
+        self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim))
+        self.register_buffer("mask_token", torch.empty(1, 1, dim))
+        self.transformer_blocks = nn.ModuleList(
+            [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps) for _ in range(num_layers)]
+        )
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=True, eps=eps)
+        self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, frames, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(batch_size, frames * height * width, channels)
+        hidden_states = self.x_embedder(hidden_states)
+        num_patches = hidden_states.shape[1]
+        hidden_states = torch.cat(
+            (
+                hidden_states,
+                self.register_tokens.expand(batch_size, -1, -1),
+                torch.zeros_like(hidden_states[:, :1]),
+            ),
+            dim=1,
+        )
+
+        position_ids = create_token_ids((frames, height, width), hidden_states.device)
+        position_ids = position_ids.expand(batch_size, -1, -1)
+        suffix_ids = position_ids.new_zeros((batch_size, self.num_register_tokens + 1, 3))
+        rotary_emb = self.pos_embed(torch.cat((position_ids, suffix_ids), dim=1))
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, rotary_emb)
+
+        hidden_states = self.proj_out(self.norm_out(hidden_states))[:, :num_patches]
+        hidden_states = hidden_states.view(
+            batch_size,
+            frames,
+            height,
+            width,
+            self.out_channels,
+            self.patch_size_t,
+            self.patch_size,
+            self.patch_size,
+        )
+        hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
+        return hidden_states.reshape(
+            batch_size,
+            self.out_channels,
+            frames * self.patch_size_t,
+            height * self.patch_size,
+            width * self.patch_size,
+        )
+
+
+class MiniMaxH3VideoVAE(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        out_ch=3,
+        ch=128,
+        embed_dim=24,
+        z_channels=24,
+        ch_mult=(1, 2, 2, 4, 4, 8),
+        num_res_blocks=2,
+        space_down=(2, 2, 2, 2, 1, 1),
+        time_down=(1, 2, 2, 1, 1, 1),
+        clip_length=17,
+        token_drop=3,
+        tile_size=256,
+        tile_overlap_min=64,
+        tiling=True,
+    ):
+        super().__init__()
+        self.vae_ratio = math.prod(space_down)
+        self.vae_ratio_t = math.prod(time_down)
+        self.clip_length = clip_length
+        self.token_drop = token_drop
+        self.frame_pre_padding = (-clip_length) % self.vae_ratio_t
+        self.tokens_chunk_size = math.ceil(clip_length / self.vae_ratio_t)
+        self.token_overlap = (-token_drop) % self.tokens_chunk_size
+        self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0)
+        self.use_tiling = tiling
+        self.tile_size = tile_size
+        self.tile_overlap_min = tile_overlap_min
+
+        self.encoder = EncoderFCN3D(
+            ch=ch,
+            ch_mult=list(ch_mult),
+            space_down=list(space_down),
+            time_down=list(time_down),
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            z_channels=z_channels,
+            double_z=True,
+        )
+        self.quant_conv = nn.Conv3d(2 * z_channels, 2 * embed_dim, kernel_size=1)
+        self.post_quant_conv = nn.Conv3d(embed_dim, z_channels, kernel_size=1)
+        self.decoder = ViT3DDecoder(
+            patch_size=self.vae_ratio,
+            patch_size_t=self.vae_ratio_t,
+            in_channels=z_channels,
+            out_channels=out_ch,
+        )
+        self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN))
+        self.register_buffer("latents_std", torch.tensor(LATENTS_STD))
+        self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1, 1), persistent=False)
+        self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False)
+
+    def _split_tiles(self, length: int) -> tuple[list[int], list[int], list[int]]:
+        if self.tile_size >= length:
+            return [0], [length], []
+        num_tiles = math.ceil(length / self.tile_size)
+        while self.tile_size * num_tiles - self.tile_overlap_min * (num_tiles - 1) < length:
+            num_tiles += 1
+        overlaps = [self.tile_overlap_min] * (num_tiles - 1)
+        remaining = self.tile_size * num_tiles - sum(overlaps) - length
+        for index in range(remaining // self.vae_ratio):
+            overlaps[index % (num_tiles - 1)] += self.vae_ratio
+        starts = [0]
+        for index in range(num_tiles - 1):
+            starts.append(starts[-1] + self.tile_size - overlaps[index])
+        return starts, [self.tile_size] * num_tiles, overlaps
+
+    @staticmethod
+    def _blend(first: torch.Tensor, second: torch.Tensor, extent: int, dim: int) -> torch.Tensor:
+        extent = min(first.shape[dim], second.shape[dim], extent)
+        positions = torch.arange(extent, device=second.device, dtype=second.dtype)
+        shape = [1] * first.ndim
+        shape[dim] = extent
+        first_weight = (1 - positions / extent).view(shape)
+        second_weight = (positions / extent).view(shape)
+        first_slice = [slice(None)] * first.ndim
+        first_slice[dim] = slice(-extent, None)
+        second_slice = [slice(None)] * second.ndim
+        second_slice[dim] = slice(0, extent)
+        blended = first[tuple(first_slice)] * first_weight + second[tuple(second_slice)] * second_weight
+        if extent == second.shape[dim]:
+            return blended
+        rest = [slice(None)] * second.ndim
+        rest[dim] = slice(extent, None)
+        return torch.cat((blended, second[tuple(rest)]), dim=dim)
+
+    def _stitch_tiles(
+        self,
+        tiles: list[list[torch.Tensor]],
+        height_overlaps: list[int],
+        width_overlaps: list[int],
+    ) -> torch.Tensor:
+        stitched_rows = []
+        for row_index, row in enumerate(tiles):
+            stitched_row = []
+            for column_index, tile in enumerate(row):
+                if row_index:
+                    tile = self._blend(tiles[row_index - 1][column_index], tile, height_overlaps[row_index - 1], -2)
+                if column_index:
+                    tile = self._blend(row[column_index - 1], tile, width_overlaps[column_index - 1], -1)
+                if row_index < len(tiles) - 1:
+                    tile = tile[..., : -height_overlaps[row_index], :]
+                if column_index < len(row) - 1:
+                    tile = tile[..., :, : -width_overlaps[column_index]]
+                stitched_row.append(tile)
+            stitched_rows.append(torch.cat(stitched_row, dim=-1))
+        return torch.cat(stitched_rows, dim=-2)
+
+    def _encode_clip(self, pixels: torch.Tensor) -> torch.Tensor:
+        if not self.use_tiling:
+            return self.quant_conv(self.encoder(pixels))
+        height_starts, height_lengths, height_overlaps = self._split_tiles(pixels.shape[-2])
+        width_starts, width_lengths, width_overlaps = self._split_tiles(pixels.shape[-1])
+        tiles = [
+            [
+                self.quant_conv(
+                    self.encoder(
+                        pixels[
+                            ...,
+                            top : top + tile_height,
+                            left : left + tile_width,
+                        ]
+                    )
+                )
+                for left, tile_width in zip(width_starts, width_lengths)
+            ]
+            for top, tile_height in zip(height_starts, height_lengths)
+        ]
+        latent_height_overlaps = [value // self.vae_ratio for value in height_overlaps]
+        latent_width_overlaps = [value // self.vae_ratio for value in width_overlaps]
+        return self._stitch_tiles(tiles, latent_height_overlaps, latent_width_overlaps)
+
+    def _decode_clip(self, latents: torch.Tensor) -> torch.Tensor:
+        if not self.use_tiling:
+            return self.decoder(self.post_quant_conv(latents))
+        height_starts, height_lengths, height_overlaps = self._split_tiles(latents.shape[-2] * self.vae_ratio)
+        width_starts, width_lengths, width_overlaps = self._split_tiles(latents.shape[-1] * self.vae_ratio)
+        tiles = [
+            [
+                self.decoder(
+                    self.post_quant_conv(
+                        latents[
+                            ...,
+                            top // self.vae_ratio : (top + tile_height) // self.vae_ratio,
+                            left // self.vae_ratio : (left + tile_width) // self.vae_ratio,
+                        ]
+                    )
+                )
+                for left, tile_width in zip(width_starts, width_lengths)
+            ]
+            for top, tile_height in zip(height_starts, height_lengths)
+        ]
+        return self._stitch_tiles(tiles, height_overlaps, width_overlaps)
+
+    def _encode_video(self, pixels: torch.Tensor) -> torch.Tensor:
+        frame_count = pixels.shape[2]
+        if frame_count % self.clip_length:
+            padding = pixels[:, :, -1:].repeat(1, 1, (-frame_count) % self.clip_length, 1, 1)
+            pixels = torch.cat((pixels, padding), dim=2)
+        moments = torch.cat(
+            [
+                self._encode_clip(pixels[:, :, start : start + self.clip_length])
+                for start in range(0, pixels.shape[2], self.clip_length)
+            ],
+            dim=2,
+        )
+        return moments[:, :, : -self.token_drop] if self.token_drop else moments
+
+    def _decode_video(self, latents: torch.Tensor) -> torch.Tensor:
+        chunk_tokens = self.tokens_chunk_size
+        chunk_frames = chunk_tokens * self.vae_ratio_t
+        padded_token_count = latents.shape[2] + self.token_drop
+        pad_tokens = (-padded_token_count) % chunk_tokens
+        num_chunks = (padded_token_count + pad_tokens) // chunk_tokens - int(self.token_drop > 0)
+        if num_chunks < 1:
+            pad_tokens += chunk_tokens
+            num_chunks = 1
+        if pad_tokens:
+            latents = torch.cat((latents, latents[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)), dim=2)
+
+        decoded = []
+        overlap = None
+        for index in range(num_chunks):
+            start = index * chunk_tokens
+            clip = self._decode_clip(latents[:, :, start : start + chunk_tokens + self.token_overlap])
+            for part_index in range(int(self.token_drop > 0) + 1):
+                frame_start = part_index * chunk_frames
+                part = clip[:, :, frame_start : frame_start + chunk_frames]
+                part = part[:, :, self.frame_pre_padding :]
+                if part_index == 0:
+                    if overlap is not None:
+                        part = self._blend(overlap, part, self.frame_overlap, dim=-3)
+                    decoded.append(part)
+                else:
+                    overlap = part
+        if overlap is not None:
+            decoded.append(overlap)
+        pixels = torch.cat(decoded, dim=2)
+
+        if pad_tokens:
+            intra_tail = self.clip_length % self.vae_ratio_t
+            tokens_before_padding = latents.shape[2] - pad_tokens
+            pad_frames = sum(
+                intra_tail if intra_tail and (tokens_before_padding + offset) % chunk_tokens == 0 else self.vae_ratio_t
+                for offset in range(pad_tokens)
+            )
+            pixels = pixels[:, :, :-pad_frames]
+        return pixels
+
+    def encode_moments(self, pixels: torch.Tensor) -> torch.Tensor:
+        if pixels.ndim == 4:
+            pixels = pixels.unsqueeze(2)
+        pixels = (pixels + 1.0) * 0.5
+        pixels = (pixels - self.pixel_mean.to(pixels)) / self.pixel_std.to(pixels)
+        if pixels.shape[2] == 1:
+            moments = self._encode_clip(pixels)[:, :, -1:]
+        else:
+            moments = self._encode_video(pixels)
+        return moments.float()
+
+    def encode(self, pixels: torch.Tensor) -> torch.Tensor:
+        mean = self.encode_moments(pixels).chunk(2, dim=1)[0]
+        latent_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(mean)
+        latent_std = self.latents_std.view(1, -1, 1, 1, 1).to(mean)
+        return (mean - latent_mean) / latent_std
+
+    def encode_tiled(self, pixels: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return self.encode(pixels)
+
+    def decode_tiled(self, latents: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return self.decode(latents)
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        latent_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(latents)
+        latent_std = self.latents_std.view(1, -1, 1, 1, 1).to(latents)
+        latents = latents * latent_std + latent_mean
+        if latents.shape[2] == 1:
+            pixels = self._decode_clip(latents)[:, :, -1:]
+        else:
+            pixels = self._decode_video(latents)
+        pixels = pixels.float() * self.pixel_std.to(pixels) + self.pixel_mean.to(pixels)
+        return pixels.clamp(0.0, 1.0) * 2.0 - 1.0
+
+
+def _video_posterior_sample(vae, pixels: torch.Tensor, generator: torch.Generator, fp16_roundtrip: bool) -> torch.Tensor:
+    moments = vae.encode_moments(pixels)
+    if moments.ndim != 5 or moments.shape[1] != 48:
+        raise ValueError(f"Expected MiniMax-H3 video moments [B,48,T,H,W], got {tuple(moments.shape)}")
+    mean, logvar = moments.float().chunk(2, dim=1)
+    noise = torch.randn(mean.shape, generator=generator, dtype=torch.float32, device="cpu").to(mean.device)
+    sample = mean + torch.exp(0.5 * logvar.clamp(-30.0, 20.0)) * noise
+    if fp16_roundtrip:
+        sample = sample.to(torch.float16).to(torch.float32)
+    latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(sample)
+    latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(sample)
+    return (sample - latents_mean) / latents_std
+
+
+@torch.no_grad()
+def encode_video_target(vae, pixels: torch.Tensor, cache_seed: int, canonical_item_key: str) -> torch.Tensor:
+    digest = hashlib.sha256(f"{cache_seed}\0{canonical_item_key}".encode()).digest()
+    seed = int.from_bytes(digest[:8], "little") % (2**63)
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    return _video_posterior_sample(vae, pixels, generator, fp16_roundtrip=False)
+
+
+@torch.no_grad()
+def encode_video_condition(vae, pixels: torch.Tensor) -> torch.Tensor:
+    generator = torch.Generator(device="cpu").manual_seed(42)
+    return _video_posterior_sample(vae, pixels, generator, fp16_roundtrip=True)
+
+
+def load_video_vae(
+    path,
+    device: str | torch.device = "cpu",
+    dtype: torch.dtype = torch.float16,
+    disable_mmap: bool = False,
+) -> MiniMaxH3VideoVAE:
+    from musubi_tuner.minimax_h3.checkpoint import load_safetensors_module, resolve_safetensors_files
+
+    files = resolve_safetensors_files(path)
+    return load_safetensors_module(
+        MiniMaxH3VideoVAE,
+        files,
+        device=device,
+        dtype=dtype,
+        key_prefixes=("first_stage_model.", "video_vae.", "vae."),
+        disable_mmap=disable_mmap,
+    )
diff --git a/musubi_tuner/modules/__init__.py b/musubi_tuner/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/musubi_tuner/modules/adafactor_fused.py b/musubi_tuner/modules/adafactor_fused.py
new file mode 100644
index 0000000000000000000000000000000000000000..79ea1c2f7fca5df9c3c6993952db84c847cd02fb
--- /dev/null
+++ b/musubi_tuner/modules/adafactor_fused.py
@@ -0,0 +1,141 @@
+# copy from sd-scripts
+
+import math
+import torch
+from transformers import Adafactor
+
+# stochastic rounding for bfloat16
+# The implementation was provided by 2kpr. Thank you very much!
+
+
+def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
+    """
+    copies source into target using stochastic rounding
+
+    Args:
+        target: the target tensor with dtype=bfloat16
+        source: the target tensor with dtype=float32
+    """
+    # create a random 16 bit integer
+    result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
+
+    # add the random number to the lower 16 bit of the mantissa
+    result.add_(source.view(dtype=torch.int32))
+
+    # mask off the lower 16 bit of the mantissa
+    result.bitwise_and_(-65536)  # -65536 = FFFF0000 as a signed int32
+
+    # copy the higher 16 bit into the target tensor
+    target.copy_(result.view(dtype=torch.float32))
+
+    del result
+
+
+@torch.no_grad()
+def adafactor_step_param(self, p, group):
+    if p.grad is None:
+        return
+    grad = p.grad
+    if grad.dtype in {torch.float16, torch.bfloat16}:
+        grad = grad.float()
+    if grad.is_sparse:
+        raise RuntimeError("Adafactor does not support sparse gradients.")
+
+    state = self.state[p]
+    grad_shape = grad.shape
+
+    factored, use_first_moment = Adafactor._get_options(group, grad_shape)
+    # State Initialization
+    if len(state) == 0:
+        state["step"] = 0
+
+        if use_first_moment:
+            # Exponential moving average of gradient values
+            state["exp_avg"] = torch.zeros_like(grad)
+        if factored:
+            state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+            state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+        else:
+            state["exp_avg_sq"] = torch.zeros_like(grad)
+
+        state["RMS"] = 0
+    else:
+        if use_first_moment:
+            state["exp_avg"] = state["exp_avg"].to(grad)
+        if factored:
+            state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+            state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+        else:
+            state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+    p_data_fp32 = p
+    if p.dtype in {torch.float16, torch.bfloat16}:
+        p_data_fp32 = p_data_fp32.float()
+
+    state["step"] += 1
+    state["RMS"] = Adafactor._rms(p_data_fp32)
+    lr = Adafactor._get_lr(group, state)
+
+    beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+    update = (grad**2) + group["eps"][0]
+    if factored:
+        exp_avg_sq_row = state["exp_avg_sq_row"]
+        exp_avg_sq_col = state["exp_avg_sq_col"]
+
+        exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
+        exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))
+
+        # Approximation of exponential moving average of square of gradient
+        update = Adafactor._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+        update.mul_(grad)
+    else:
+        exp_avg_sq = state["exp_avg_sq"]
+
+        exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
+        update = exp_avg_sq.rsqrt().mul_(grad)
+
+    update.div_((Adafactor._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+    update.mul_(lr)
+
+    if use_first_moment:
+        exp_avg = state["exp_avg"]
+        exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"]))
+        update = exp_avg
+
+    if group["weight_decay"] != 0:
+        p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr))
+
+    p_data_fp32.add_(-update)
+
+    # if p.dtype in {torch.float16, torch.bfloat16}:
+    #    p.copy_(p_data_fp32)
+
+    if p.dtype == torch.bfloat16:
+        copy_stochastic_(p, p_data_fp32)
+    elif p.dtype == torch.float16:
+        p.copy_(p_data_fp32)
+
+
+@torch.no_grad()
+def adafactor_step(self, closure=None):
+    """
+    Performs a single optimization step
+
+    Arguments:
+        closure (callable, optional): A closure that reevaluates the model
+            and returns the loss.
+    """
+    loss = None
+    if closure is not None:
+        loss = closure()
+
+    for group in self.param_groups:
+        for p in group["params"]:
+            adafactor_step_param(self, p, group)
+
+    return loss
+
+
+def patch_adafactor_fused(optimizer: Adafactor):
+    optimizer.step_param = adafactor_step_param.__get__(optimizer)
+    optimizer.step = adafactor_step.__get__(optimizer)
diff --git a/musubi_tuner/modules/attention.py b/musubi_tuner/modules/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..236314c085b89307b1c4f0e5b76c18b30a505675
--- /dev/null
+++ b/musubi_tuner/modules/attention.py
@@ -0,0 +1,297 @@
+# Unified attention function supporting various implementations
+
+from dataclasses import dataclass
+import torch
+from typing import Optional, Union
+
+try:
+    import flash_attn
+    from flash_attn.flash_attn_interface import _flash_attn_forward
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+    from flash_attn.flash_attn_interface import flash_attn_func
+except ImportError:
+    flash_attn = None
+    flash_attn_varlen_func = None
+    _flash_attn_forward = None
+    flash_attn_func = None
+
+try:
+    from sageattention import sageattn_varlen, sageattn
+except ImportError:
+    sageattn_varlen = None
+    sageattn = None
+
+try:
+    import xformers.ops as xops
+except ImportError:
+    xops = None
+
+
+@dataclass
+class AttentionParams:
+    attn_mode: Optional[str] = None
+    split_attn: bool = False
+    img_len: Optional[int] = None
+    attention_mask: Optional[torch.Tensor] = None
+    seqlens: Optional[torch.Tensor] = None
+    cu_seqlens: Optional[torch.Tensor] = None
+    max_seqlen: Optional[int] = None
+
+    @staticmethod
+    def create_attention_params(attn_mode: Optional[str], split_attn: bool) -> "AttentionParams":
+        return AttentionParams(attn_mode, split_attn)
+
+    @staticmethod
+    def create_attention_params_from_mask(
+        attn_mode: Optional[str], split_attn: bool, img_len: Optional[int], attention_mask: Optional[torch.Tensor]
+    ) -> "AttentionParams":
+        if attention_mask is None:
+            # No attention mask provided: assume all tokens are valid
+            return AttentionParams(attn_mode, split_attn, None, None, None, None, None)
+        else:
+            # Note: attention_mask is only for text tokens, not including image tokens
+            seqlens = attention_mask.sum(dim=1).to(torch.int32) + img_len  # [B]
+            max_seqlen = attention_mask.shape[1] + img_len
+
+            if split_attn:
+                # cu_seqlens is not needed for split attention
+                return AttentionParams(attn_mode, split_attn, img_len, attention_mask, seqlens, None, max_seqlen)
+
+            # Convert attention mask to cumulative sequence lengths for flash attention
+            batch_size = attention_mask.shape[0]
+            cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device=attention_mask.device)
+            for i in range(batch_size):
+                cu_seqlens[2 * i + 1] = i * max_seqlen + seqlens[i]  # end of valid tokens for query
+                cu_seqlens[2 * i + 2] = (i + 1) * max_seqlen  # end of all tokens for query
+
+            # Expand attention mask to include image tokens
+            attention_mask = torch.nn.functional.pad(attention_mask, (img_len, 0), value=1)  # [B, img_len + L]
+
+            # attention bias for xformers
+            if attn_mode == "xformers":
+                seqlens_list = seqlens.cpu().tolist()
+                attention_mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
+                    seqlens_list, seqlens_list, device=attention_mask.device
+                )
+            elif attn_mode == "torch":
+                attention_mask = attention_mask[:, None, None, :].to(torch.bool)  # [B, 1, 1, img_len + L]
+
+            return AttentionParams(attn_mode, split_attn, img_len, attention_mask, seqlens, cu_seqlens, max_seqlen)
+
+
+def attention(
+    qkv_or_q: Union[torch.Tensor, list],
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    attn_params: Optional[AttentionParams] = None,
+    drop_rate: float = 0.0,
+) -> torch.Tensor:
+    """
+    Compute scaled dot-product attention with variable sequence lengths.
+
+    Handles batches with different sequence lengths by splitting and
+    processing each sequence individually.
+
+    Args:
+        qkv_or_q: Query tensor [B, L, H, D]. or list of such tensors.
+        k: Key tensor [B, L, H, D].
+        v: Value tensor [B, L, H, D].
+        attn_param: Attention parameters including mask and sequence lengths.
+        drop_rate: Attention dropout rate.
+
+    Returns:
+        Attention output tensor [B, L, H*D].
+    """
+    if isinstance(qkv_or_q, list):
+        q, k, v = qkv_or_q
+        q: torch.Tensor = q
+        qkv_or_q.clear()
+        del qkv_or_q
+    else:
+        q: torch.Tensor = qkv_or_q
+        del qkv_or_q
+        assert k is not None and v is not None, "k and v must be provided if qkv_or_q is a tensor"
+    if attn_params is None:
+        attn_params = AttentionParams.create_attention_params("torch", False)
+
+    # GQA: q may carry more heads than k/v (e.g. Krea 2 = 48 query / 12 kv heads). flash and
+    # sageattn group heads natively inside the kernel (verified), so they ignore this. For the
+    # torch (SDPA) path we expand k/v to q's head count below instead of passing enable_gqa=True,
+    # because enable_gqa forces SDPA onto the slow math kernel (~7x slower than the fused kernels
+    # at K2 scale); the repeat is numerically identical. Revert to enable_gqa=True once PyTorch's
+    # fused (flash / mem-efficient) kernels gain native GQA support. (q/k/v here are [B, L, H, D].)
+    enable_gqa = q.shape[-2] != k.shape[-2]
+
+    # If split attn is False, attention mask is provided and all sequence lengths are same, we can trim the sequence
+    seqlen_trimmed = False
+    # Trim if all seqlens are the same, for attention modes other than flash or sageattn (which can handle masks efficiently)
+    if (
+        not attn_params.split_attn
+        and attn_params.attention_mask is not None
+        and attn_params.seqlens is not None
+        and (attn_params.attn_mode != "flash" and attn_params.attn_mode != "sageattn")
+    ):
+        if torch.all(attn_params.seqlens == attn_params.seqlens[0]):
+            seqlen = attn_params.seqlens[0].item()
+            q = q[:, :seqlen]
+            k = k[:, :seqlen]
+            v = v[:, :seqlen]
+            max_seqlen = attn_params.max_seqlen
+            attn_params = AttentionParams.create_attention_params(attn_params.attn_mode, False)  # do not in-place modify
+            attn_params.max_seqlen = max_seqlen  # keep max_seqlen for padding
+            seqlen_trimmed = True
+
+    # Determine tensor layout based on attention implementation
+    if attn_params.attn_mode == "torch" or (
+        attn_params.attn_mode == "sageattn" and (attn_params.split_attn or attn_params.cu_seqlens is None)
+    ):
+        transpose_fn = lambda x: x.transpose(1, 2)  # [B, H, L, D] for SDPA and sageattn with fixed length
+        # pad on sequence length dimension
+        pad_fn = lambda x, pad_to: torch.nn.functional.pad(x, (0, 0, 0, pad_to - x.shape[-2]), value=0)
+    else:
+        transpose_fn = lambda x: x  # [B, L, H, D] for other implementations
+        # pad on sequence length dimension
+        pad_fn = lambda x, pad_to: torch.nn.functional.pad(x, (0, 0, 0, 0, 0, pad_to - x.shape[-3]), value=0)
+
+    # Process each batch element with its valid sequence lengths
+    if attn_params.split_attn:
+        if attn_params.seqlens is None:
+            # If no seqlens provided, assume all tokens are valid
+            attn_params = AttentionParams.create_attention_params(attn_params.attn_mode, True)  # do not in-place modify
+            attn_params.seqlens = torch.tensor([q.shape[1]] * q.shape[0], device=q.device)
+            attn_params.max_seqlen = q.shape[1]
+        q = [transpose_fn(q[i : i + 1, : attn_params.seqlens[i]]) for i in range(len(q))]
+        k = [transpose_fn(k[i : i + 1, : attn_params.seqlens[i]]) for i in range(len(k))]
+        v = [transpose_fn(v[i : i + 1, : attn_params.seqlens[i]]) for i in range(len(v))]
+    else:
+        q = transpose_fn(q)
+        k = transpose_fn(k)
+        v = transpose_fn(v)
+
+    if attn_params.attn_mode == "torch":
+        if attn_params.split_attn:
+            x = []
+            for i in range(len(q)):
+                qi, ki, vi = q[i], k[i], v[i]
+                if enable_gqa:  # expand k/v heads to avoid SDPA's slow enable_gqa math path
+                    g = qi.shape[1] // ki.shape[1]  # [B, H, L, D] -> heads at dim 1
+                    ki = ki.repeat_interleave(g, dim=1)
+                    vi = vi.repeat_interleave(g, dim=1)
+                x_i = torch.nn.functional.scaled_dot_product_attention(qi, ki, vi, dropout_p=drop_rate)
+                q[i] = None
+                k[i] = None
+                v[i] = None
+                x.append(pad_fn(x_i, attn_params.max_seqlen))  # B, H, L, D
+            x = torch.cat(x, dim=0)
+            q, k, v = None, None, None
+
+        else:
+            if enable_gqa:  # expand k/v heads to avoid SDPA's slow enable_gqa math path
+                g = q.shape[1] // k.shape[1]  # [B, H, L, D] -> heads at dim 1
+                k = k.repeat_interleave(g, dim=1)
+                v = v.repeat_interleave(g, dim=1)
+            x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_params.attention_mask, dropout_p=drop_rate)
+            q, k, v = None, None, None
+
+    elif attn_params.attn_mode == "xformers":
+        if attn_params.split_attn:
+            x = []
+            for i in range(len(q)):
+                qi, ki, vi = q[i], k[i], v[i]
+                if enable_gqa:  # expand k/v heads to q's count; xformers 4D has no native GQA ([B, L, H, D] -> dim 2)
+                    g = qi.shape[2] // ki.shape[2]
+                    ki = ki.repeat_interleave(g, dim=2)
+                    vi = vi.repeat_interleave(g, dim=2)
+                x_i = xops.memory_efficient_attention(qi, ki, vi, p=drop_rate)
+                q[i] = None
+                k[i] = None
+                v[i] = None
+                x.append(pad_fn(x_i, attn_params.max_seqlen))  # B, L, H, D
+            x = torch.cat(x, dim=0)
+            q, k, v = None, None, None
+
+        else:
+            if enable_gqa:  # expand k/v heads to q's count; xformers 4D has no native GQA ([B, L, H, D] -> dim 2)
+                g = q.shape[2] // k.shape[2]
+                k = k.repeat_interleave(g, dim=2)
+                v = v.repeat_interleave(g, dim=2)
+            x = xops.memory_efficient_attention(q, k, v, attn_bias=attn_params.attention_mask, p=drop_rate)
+            q, k, v = None, None, None
+
+    elif attn_params.attn_mode == "sageattn":
+        if attn_params.split_attn:
+            x = []
+            for i in range(len(q)):
+                # HND seems to cause an error
+                x_i = sageattn(q[i], k[i], v[i])  # B, H, L, D. No dropout support
+                q[i] = None
+                k[i] = None
+                v[i] = None
+                x.append(pad_fn(x_i, attn_params.max_seqlen))  # B, H, L, D
+            x = torch.cat(x, dim=0)
+            q, k, v = None, None, None
+        elif attn_params.cu_seqlens is None:  # all tokens are valid
+            x = sageattn(q, k, v)  # B, L, H, D. No dropout support
+            q, k, v = None, None, None
+        else:
+            # Reshape to [(bxs), a, d]
+            batch_size, seqlen = q.shape[0], q.shape[1]
+            # reshape (not view): callers may pass non-contiguous q/k/v (e.g. Krea 2 transposes
+            # [B,H,L,D]->[B,L,H,D] after RoPE); reshape copies only when needed, view-equivalent otherwise.
+            q = q.reshape(q.shape[0] * q.shape[1], *q.shape[2:])  # [B*L, H, D]
+            k = k.reshape(k.shape[0] * k.shape[1], *k.shape[2:])  # [B*L, H, D]
+            v = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:])  # [B*L, H, D]
+
+            # Assume cu_seqlens_q == cu_seqlens_kv and max_seqlen_q == max_seqlen_kv. No dropout support
+            x = sageattn_varlen(
+                q, k, v, attn_params.cu_seqlens, attn_params.cu_seqlens, attn_params.max_seqlen, attn_params.max_seqlen
+            )
+            q, k, v = None, None, None
+
+            # Reshape x with shape [(bxs), a, d] to [b, s, a, d]
+            x = x.view(batch_size, seqlen, x.shape[-2], x.shape[-1])  # B, L, H, D
+
+    elif attn_params.attn_mode == "flash":
+        if attn_params.split_attn:
+            x = []
+            for i in range(len(q)):
+                # HND seems to cause an error
+                x_i = flash_attn_func(q[i], k[i], v[i], drop_rate)  # B, L, H, D
+                q[i] = None
+                k[i] = None
+                v[i] = None
+                x.append(pad_fn(x_i, attn_params.max_seqlen))  # B, L, H, D
+            x = torch.cat(x, dim=0)
+            q, k, v = None, None, None
+        elif attn_params.cu_seqlens is None:  # all tokens are valid
+            x = flash_attn_func(q, k, v, drop_rate)  # B, L, H, D
+            q, k, v = None, None, None
+        else:
+            # Reshape to [(bxs), a, d]
+            batch_size, seqlen = q.shape[0], q.shape[1]
+            # reshape (not view): callers may pass non-contiguous q/k/v (e.g. Krea 2 transposes
+            # [B,H,L,D]->[B,L,H,D] after RoPE); reshape copies only when needed, view-equivalent otherwise.
+            q = q.reshape(q.shape[0] * q.shape[1], *q.shape[2:])  # [B*L, H, D]
+            k = k.reshape(k.shape[0] * k.shape[1], *k.shape[2:])  # [B*L, H, D]
+            v = v.reshape(v.shape[0] * v.shape[1], *v.shape[2:])  # [B*L, H, D]
+
+            # Assume cu_seqlens_q == cu_seqlens_kv and max_seqlen_q == max_seqlen_kv
+            x = flash_attn_varlen_func(
+                q, k, v, attn_params.cu_seqlens, attn_params.cu_seqlens, attn_params.max_seqlen, attn_params.max_seqlen, drop_rate
+            )
+            q, k, v = None, None, None
+
+            # Reshape x with shape [(bxs), a, d] to [b, s, a, d]
+            x = x.view(batch_size, seqlen, x.shape[-2], x.shape[-1])  # B, L, H, D
+
+    else:
+        # Currently only PyTorch SDPA and xformers are implemented
+        raise ValueError(f"Unsupported attention mode: {attn_params.attn_mode}")
+
+    x = transpose_fn(x)  # [B, L, H, D]
+    x = x.reshape(x.shape[0], x.shape[1], -1)  # [B, L, H*D]
+
+    if seqlen_trimmed:
+        x = torch.nn.functional.pad(x, (0, 0, 0, attn_params.max_seqlen - x.shape[1]), value=0)  # pad back to max_seqlen
+
+    return x
diff --git a/musubi_tuner/modules/custom_offloading_utils.py b/musubi_tuner/modules/custom_offloading_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ae309162a5edd52dfca2c0659ac14114cd8c1e3
--- /dev/null
+++ b/musubi_tuner/modules/custom_offloading_utils.py
@@ -0,0 +1,1096 @@
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from dataclasses import dataclass
+import gc
+import time
+from typing import Optional
+import torch
+import torch.nn as nn
+
+
+# Keep these functions here for portability, and private to avoid confusion with the ones in device_utils.py
+def _clean_memory_on_device(device: torch.device):
+    r"""
+    Clean memory on the specified device, will be called from training scripts.
+    """
+    gc.collect()
+
+    # device may "cuda" or "cuda:0", so we need to check the type of device
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    if device.type == "xpu":
+        torch.xpu.empty_cache()
+    if device.type == "mps":
+        torch.mps.empty_cache()
+
+
+def _synchronize_device(device: torch.device):
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    elif device.type == "xpu":
+        torch.xpu.synchronize()
+    elif device.type == "mps":
+        torch.mps.synchronize()
+
+
+def swap_weight_devices_no_cuda(device: torch.device, layer_to_cpu: nn.Module, layer_to_cuda: nn.Module):
+    """
+    not tested
+    """
+    assert layer_to_cpu.__class__ == layer_to_cuda.__class__
+
+    weight_swap_jobs = []
+    for module_to_cpu, module_to_cuda in zip(layer_to_cpu.modules(), layer_to_cuda.modules()):
+        if hasattr(module_to_cpu, "weight") and module_to_cpu.weight is not None:
+            weight_swap_jobs.append((module_to_cpu, module_to_cuda, module_to_cpu.weight.data, module_to_cuda.weight.data))
+
+    # device to cpu
+    for module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view in weight_swap_jobs:
+        module_to_cpu.weight.data = cuda_data_view.data.to("cpu", non_blocking=True)
+
+    _synchronize_device(device)
+
+    # cpu to device
+    for module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view in weight_swap_jobs:
+        cuda_data_view.copy_(module_to_cuda.weight.data, non_blocking=True)
+        module_to_cuda.weight.data = cuda_data_view
+
+    _synchronize_device(device)
+
+
+def weighs_to_device(layer: nn.Module, device: torch.device):
+    for module in layer.modules():
+        if hasattr(module, "weight") and module.weight is not None and module.__class__.__name__.endswith("Linear"):
+            module.weight.data = module.weight.data.to(device, non_blocking=device.type != "cpu")
+
+
+@dataclass
+class BlockSwapConfig:
+    """
+    Construction policy for a block-swap offloader, assembled once by the training/inference script and
+    passed through each architecture's ``enable_block_swap`` to ``create_offloader``.
+
+    It holds everything the offloader constructor needs that is *not* architecture-specific. The
+    architecture-specific arguments (the block-type label, the block list, and its counts) are supplied by
+    ``enable_block_swap`` per block list. Adding a new offloader knob (or a whole new offloader type) means
+    adding a field here and a branch in ``create_offloader`` -- the per-architecture ``enable_block_swap``
+    signatures stay ``(blocks_to_swap, config)`` and never change.
+    """
+
+    device: torch.device
+    supports_backward: bool
+    use_pinned_memory: bool = False
+    h2d_only: bool = False  # frozen-base (LoRA / LoHa / LoKr) only: H2D-only streaming, no device->host copy
+    ring_size: int = 2  # (h2d_only) number of GPU ring buffers for streamed blocks; 2 = double buffering
+    debug: bool = False
+
+    @classmethod
+    def from_args(cls, args, device: torch.device, supports_backward: bool) -> "BlockSwapConfig":
+        """Build from a parsed-args namespace, tolerating scripts whose parser lacks the optional knobs."""
+        h2d_only = getattr(args, "block_swap_h2d_only", False)
+
+        # H2D-only streams frozen weights into a reused GPU ring buffer and loads them with an in-place copy_,
+        # which bumps the autograd version of the saved weight view. Gradient checkpointing re-reads the weight
+        # at recompute time (a short, safe window), so the version stays consistent; without it the backward
+        # pass fails with an opaque "modified by an inplace operation ... version N expected M" error that does
+        # not point at the offloader. Fail early with an actionable message instead. (Inference / forward-only
+        # has no backward, so it is unaffected.)
+        if h2d_only and supports_backward and not getattr(args, "gradient_checkpointing", False):
+            raise ValueError(
+                "--block_swap_h2d_only requires --gradient_checkpointing for training. H2D-only block swap streams"
+                " frozen weights through a reused GPU ring buffer, which advances the autograd version of weights"
+                " saved for backward; gradient checkpointing re-reads them at recompute time and avoids this."
+                " / --block_swap_h2d_only は学習時に --gradient_checkpointing が必須です（リングバッファの上書きで"
+                "backward 用に保存された重みの version が進むため。gradient checkpointing は再計算時に読み直すので回避できます）。"
+            )
+
+        ring_size = getattr(args, "block_swap_ring_size", 2)
+        if ring_size < 1:
+            raise ValueError("--block_swap_ring_size must be >= 1")
+
+        return cls(
+            device=device,
+            supports_backward=supports_backward,
+            use_pinned_memory=getattr(args, "use_pinned_memory_for_block_swap", False),
+            h2d_only=h2d_only,
+            ring_size=ring_size,
+        )
+
+
+def create_offloader(block_type: str, blocks: list[nn.Module], num_blocks: int, blocks_to_swap: int, config: BlockSwapConfig):
+    """
+    Create the block-swap offloader for one block list. This is the single place that selects the offloader
+    implementation from ``config``; ``enable_block_swap`` only supplies the per-block-list arguments, so a new
+    offloader type plugs in here without touching any architecture.
+    """
+    if config.h2d_only:
+        # H2D-only streaming for frozen-base (LoRA) training: keep a CPU master, copy Host->Device only.
+        return LoRAStreamOffloader(
+            block_type,
+            blocks,
+            num_blocks,
+            blocks_to_swap,
+            config.supports_backward,
+            config.device,
+            ring_size=config.ring_size,
+            use_pinned_memory=config.use_pinned_memory,
+            debug=config.debug,
+        )
+    return ModelOffloader(
+        block_type,
+        blocks,
+        num_blocks,
+        blocks_to_swap,
+        config.supports_backward,
+        config.device,
+        config.use_pinned_memory,
+        debug=config.debug,
+    )
+
+
+class Offloader:
+    """
+    common offloading class
+    """
+
+    def __init__(
+        self,
+        block_type: str,
+        num_blocks: int,
+        blocks_to_swap: int,
+        device: torch.device,
+        use_pinned_memory: bool = False,
+        debug: bool = False,
+    ):
+        self.block_type = block_type
+        self.num_blocks = num_blocks
+        self.blocks_to_swap = blocks_to_swap
+        self.device = device
+        self.use_pinned_memory = use_pinned_memory
+
+        # check if debug is enabled from os environment variable
+        if not debug:
+            import os
+
+            debug = os.getenv("MUSUBI_TUNER_OFFLOADER_DEBUG", "0") == "1"
+
+        self.debug = debug
+        self.debug_block_count = 0
+
+        self.thread_pool = ThreadPoolExecutor(max_workers=1)
+        self.futures = {}
+        self.cuda_available = device.type == "cuda"
+        self.stream = torch.cuda.Stream(device=device) if self.cuda_available else None
+
+        # Staging buffers for cuda offloading without large pinned memory. These are pinned memory buffers to speed up the transfer between CPU and GPU
+        # We create one staging buffer per transfer direction (A: GPU to CPU, B: CPU to GPU)
+        self.staging_buffer_a = None
+        self.staging_buffer_b = None
+
+        # Pinned buffer for cuda offloading with pinned memory. We need only one pinned buffer per layer transfer
+        self.pinned_buffer = None
+
+    def swap_weight_devices_cuda(self, device: torch.device, layer_to_cpu: nn.Module, layer_to_cuda: nn.Module):
+        assert layer_to_cpu.__class__ == layer_to_cuda.__class__
+
+        debug_print = False
+        if self.debug:
+            debug_print = self.debug_block_count % 10 == 0
+            self.debug_block_count += 1
+
+        class Timer:
+            def __init__(self, enabled=False):
+                self.enabled = enabled
+                self.totals = defaultdict(float)
+                self.start_time = time.perf_counter()
+
+            @contextmanager
+            def section(self, name):
+                if not self.enabled:
+                    yield
+                    return
+                t0 = time.perf_counter()
+                try:
+                    yield
+                finally:
+                    self.totals[name] += time.perf_counter() - t0
+
+        T = Timer(enabled=debug_print)
+
+        weight_swap_jobs = []
+
+        # This is not working for all cases (e.g. SD3), so we need to find the corresponding modules. kept here for reference:
+        # for module_to_cpu, module_to_cuda in zip(layer_to_cpu.modules(), layer_to_cuda.modules()):
+        #     print(module_to_cpu.__class__, module_to_cuda.__class__)
+        #     if hasattr(module_to_cpu, "weight") and module_to_cpu.weight is not None:
+        #         weight_swap_jobs.append((module_to_cpu, module_to_cuda, module_to_cpu.weight.data, module_to_cuda.weight.data))
+
+        with T.section("find modules"):
+            modules_to_cpu = {k: v for k, v in layer_to_cpu.named_modules()}
+            for module_to_cuda_name, module_to_cuda in layer_to_cuda.named_modules():
+                if (
+                    hasattr(module_to_cuda, "weight")
+                    and module_to_cuda.weight is not None
+                    and module_to_cuda.__class__.__name__.endswith("Linear")
+                ):
+                    module_to_cpu = modules_to_cpu.get(module_to_cuda_name, None)
+                    if module_to_cpu is not None and module_to_cpu.weight.shape == module_to_cuda.weight.shape:
+                        weight_swap_jobs.append(
+                            (module_to_cpu, module_to_cuda, module_to_cpu.weight.data, module_to_cuda.weight.data)
+                        )
+                    else:
+                        if module_to_cuda.weight.data.device.type != device.type:
+                            module_to_cuda.weight.data = module_to_cuda.weight.data.to(device)
+
+        with T.section("synchronize before swap"):
+            torch.cuda.current_stream().synchronize()  # this prevents the illegal loss value by ensuring offloading layer's calculation is done
+
+        if not self.use_pinned_memory:
+            # Minimize using pinned memory for lower shared GPU RAM usage
+            stream = self.stream
+            with torch.cuda.stream(stream):
+                if self.staging_buffer_a is None:
+                    # Create staging buffer as pinned memory (as shared GPU ram). We specify device for correct pinning on multi-GPU systems
+                    self.staging_buffer_a = [
+                        torch.empty_like(cuda_data_view, device="cpu").pin_memory(device=device)
+                        for _, _, cuda_data_view, _ in weight_swap_jobs
+                    ]
+                    self.staging_buffer_b = [
+                        torch.empty_like(cuda_data_view, device="cpu").pin_memory(device=device)
+                        for _, _, cuda_data_view, _ in weight_swap_jobs
+                    ]
+
+                # Copy weights to staging buffers and record events
+                event_b = None
+                for sbuf_a, sbuf_b, (module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view) in zip(
+                    self.staging_buffer_a, self.staging_buffer_b, weight_swap_jobs
+                ):
+                    # CUDA to staging buffer A, non-blocking copy
+                    event_a = torch.cuda.Event()
+                    with T.section("cuda to staging A"):
+                        sbuf_a.copy_(cuda_data_view.data, non_blocking=True)
+                        event_a.record(stream)
+
+                    # Wait for staging buffer B to be ready
+                    if event_b is not None:
+                        with T.section("wait staging B"):
+                            event_b.synchronize()  # synchronize is needed to wait CPU process. wait_event does not work here because it waits on GPU side only
+
+                    # CPU to staging buffer B, CPU to pinned CPU, synchronous copy. Can overlap with CUDA to staging buffer A
+                    with T.section("cpu to staging B"):
+                        # Making this multithreaded does not help, and 'non_blocking=True' does not help either.
+                        sbuf_b.copy_(module_to_cuda.weight.data)  # BOTTLENECK
+
+                    # Wait for staging buffer A to be ready, and CUDA data view can be reused
+                    with T.section("wait staging A"):
+                        event_a.synchronize()
+
+                    # Staging buffer B to CUDA, non-blocking copy.
+                    event_b = torch.cuda.Event()
+                    with T.section("staging B to CUDA"):
+                        cuda_data_view.copy_(sbuf_b, non_blocking=True)
+                        event_b.record(stream)
+
+                    # Staging buffer A to CPU, synchronous copy. Can overlap with staging buffer B to CUDA
+                    with T.section("staging A to CPU"):
+                        cpu_data_view.copy_(sbuf_a)  # BOTTLENECK
+
+            for sbuf_a, sbuf_b, (module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view) in zip(
+                self.staging_buffer_a, self.staging_buffer_b, weight_swap_jobs
+            ):
+                # Update references
+                module_to_cuda.weight.data = cuda_data_view
+                module_to_cpu.weight.data = cpu_data_view
+
+            sync_event = event_b  # final sync event for CPU to CUDA copy
+
+        else:
+            # Use pinned memory for faster transfer between CPU and GPU, but it requires more memory
+            if self.pinned_buffer is None:
+                with torch.cuda.stream(self.stream):
+                    # Create pinned buffer as pinned memory (as shared GPU ram). We specify device for correct pinning on multi-GPU systems
+                    self.pinned_buffer = [
+                        torch.empty_like(cuda_data_view, device="cpu").pin_memory(device=device)
+                        for _, _, cuda_data_view, _ in weight_swap_jobs
+                    ]
+                self.stream.synchronize()
+            released_pinned_buffer = []
+
+            events = [torch.cuda.Event() for _ in weight_swap_jobs]  # Waiting events for GPU to CPU non-blocking copy
+
+            # Copy weights to CPU
+            for event, module_pin_buf, (module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view) in zip(
+                events, self.pinned_buffer, weight_swap_jobs
+            ):
+                # CUDA to CPU, non-blocking copy
+                with torch.cuda.stream(self.stream):
+                    with T.section("cuda to cpu"):
+                        module_pin_buf.copy_(cuda_data_view, non_blocking=True)
+                        event.record(self.stream)
+
+            # CPU to CUDA
+            for event, (module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view) in zip(events, weight_swap_jobs):
+                with torch.cuda.stream(self.stream):
+                    # Wait for cuda_data_view to be ready
+                    with T.section("wait cpu"):
+                        self.stream.wait_event(event)
+
+                    # CPU to CUDA, non-blocking copy
+                    with T.section("cpu to cuda"):
+                        cuda_data_view.copy_(cpu_data_view, non_blocking=True)
+
+            # Update references
+            for module_pin_buf, (module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view) in zip(
+                self.pinned_buffer, weight_swap_jobs
+            ):
+                module_to_cuda.weight.data = cuda_data_view
+                module_to_cpu.weight.data = module_pin_buf
+                released_pinned_buffer.append(cpu_data_view)  # CPU data view can be reused as pinned buffer
+
+            # Reuse released pinned buffers
+            if not released_pinned_buffer[0].is_pinned():
+                # In first time, we need to create pinned buffers because offloaded weights are not pinned yet
+                with torch.cuda.stream(self.stream):
+                    released_pinned_buffer = [
+                        torch.empty_like(cuda_data_view, device="cpu").pin_memory(device=device)
+                        for _, _, cuda_data_view, _ in weight_swap_jobs
+                    ]
+            self.pinned_buffer = released_pinned_buffer
+
+            sync_event = self.stream.record_event()
+
+        if debug_print:
+            print(f"[{self.block_type}] Weight swap timing at {self.debug_block_count - 1}:")
+            for name, total in T.totals.items():
+                print(f"  {name}: {total * 1000:.2f}ms")
+            print(
+                f"Overall time: {(time.perf_counter() - T.start_time) * 1000:.2f}ms, total time in sections: {sum(T.totals.values()) * 1000:.2f}ms"
+            )
+        # print(
+        #     f"[{self.block_type}] Swapped weights in {time.perf_counter() - start_time:.2f}s. Count of modules swapped: {len(weight_swap_jobs)}"
+        # )
+
+        return sync_event
+
+    def swap_weight_devices(self, block_to_cpu: nn.Module, block_to_cuda: nn.Module):
+        if self.cuda_available:
+            sync_event = self.swap_weight_devices_cuda(self.device, block_to_cpu, block_to_cuda)
+        else:
+            swap_weight_devices_no_cuda(self.device, block_to_cpu, block_to_cuda)
+            sync_event = None
+        return sync_event
+
+    def _submit_move_blocks(self, blocks, block_idx_to_cpu, block_idx_to_cuda):
+        def move_blocks(bidx_to_cpu, block_to_cpu, bidx_to_cuda, block_to_cuda):
+            if self.debug:
+                start_time = time.perf_counter()
+                print(
+                    f"[{self.block_type}] Move block {bidx_to_cpu} to CPU and block {bidx_to_cuda} to {'CUDA' if self.cuda_available else 'device'}"
+                )
+
+            dev = self.device.index if self.device.index is not None else torch.cuda.current_device()
+            torch.cuda.set_device(dev)
+
+            sync_event = self.swap_weight_devices(block_to_cpu, block_to_cuda)
+
+            if self.debug:
+                print(
+                    f"[{self.block_type}] Moved blocks {bidx_to_cpu} to CPU and {bidx_to_cuda} to {'CUDA' if self.cuda_available else 'device'} in {time.perf_counter() - start_time:.2f}s"
+                )
+            return bidx_to_cpu, bidx_to_cuda, sync_event
+
+        block_to_cpu = blocks[block_idx_to_cpu]
+        block_to_cuda = blocks[block_idx_to_cuda]
+
+        self.futures[block_idx_to_cuda] = self.thread_pool.submit(
+            move_blocks, block_idx_to_cpu, block_to_cpu, block_idx_to_cuda, block_to_cuda
+        )
+
+    def _wait_blocks_move(self, block_idx):
+        if block_idx not in self.futures:
+            return
+
+        if self.debug:
+            print(f"[{self.block_type}] Wait for block {block_idx}")
+            start_time = time.perf_counter()
+
+        future = self.futures.pop(block_idx)
+        _, bidx_to_cuda, sync_event = future.result()
+
+        assert block_idx == bidx_to_cuda, f"Block index mismatch: {block_idx} != {bidx_to_cuda}"
+
+        if self.cuda_available and sync_event is not None:
+            # this does not wait CPU side, so the log below should be immediate when pinned memory is used
+            torch.cuda.current_stream().wait_event(sync_event)
+
+        if self.debug:
+            print(f"[{self.block_type}] Waited for block {block_idx}: {time.perf_counter() - start_time:.2f}s")
+
+
+class ModelOffloader(Offloader):
+    """
+    supports forward offloading
+    """
+
+    def __init__(
+        self,
+        block_type: str,
+        blocks: list[nn.Module],
+        num_blocks: int,
+        blocks_to_swap: int,
+        supports_backward: bool,
+        device: torch.device,
+        use_pinned_memory: bool = False,
+        debug: bool = False,
+    ):
+        super().__init__(block_type, num_blocks, blocks_to_swap, device, use_pinned_memory, debug)
+
+        self.supports_backward = supports_backward
+        self.forward_only = not supports_backward  # forward only offloading: can be changed to True for inference
+
+        if self.supports_backward:
+            # register backward hooks
+            self.remove_handles = []
+            for i, block in enumerate(blocks):
+                hook = self.create_backward_hook(blocks, i)
+                if hook is not None:
+                    handle = block.register_full_backward_hook(hook)
+                    self.remove_handles.append(handle)
+
+    def set_forward_only(self, forward_only: bool):
+        # switching must wait for all pending transfers
+        for block_idx in list(self.futures.keys()):
+            self._wait_blocks_move(block_idx)
+
+        self.forward_only = forward_only
+
+    def __del__(self):
+        if self.supports_backward:
+            for handle in self.remove_handles:
+                handle.remove()
+
+    def create_backward_hook(self, blocks: list[nn.Module], block_index: int) -> Optional[callable]:
+        # -1 for 0-based index
+        num_blocks_propagated = self.num_blocks - block_index - 1
+        swapping = num_blocks_propagated > 0 and num_blocks_propagated <= self.blocks_to_swap
+        waiting = block_index > 0 and block_index <= self.blocks_to_swap
+
+        if not swapping and not waiting:
+            return None
+
+        # create  hook
+        block_idx_to_cpu = self.num_blocks - num_blocks_propagated
+        block_idx_to_cuda = self.blocks_to_swap - num_blocks_propagated
+        block_idx_to_wait = block_index - 1
+
+        def backward_hook(module, grad_input, grad_output):
+            if self.debug:
+                print(f"Backward hook for block {block_index}")
+
+            if swapping:
+                self._submit_move_blocks(blocks, block_idx_to_cpu, block_idx_to_cuda)
+            if waiting:
+                self._wait_blocks_move(block_idx_to_wait)
+            return None
+
+        return backward_hook
+
+    def prepare_block_devices_before_forward(self, blocks: list[nn.Module]):
+        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
+            return
+
+        if self.debug:
+            print(f"[{self.block_type}] Prepare block devices before forward")
+
+        for b in blocks[0 : self.num_blocks - self.blocks_to_swap]:
+            b.to(self.device)
+            weighs_to_device(b, self.device)  # make sure weights are on device
+
+        cpu_device = torch.device("cpu")
+        for b in blocks[self.num_blocks - self.blocks_to_swap :]:
+            b.to(self.device)  # move block to device first. this makes sure that buffers (non weights) are on the device
+            weighs_to_device(b, cpu_device)  # make sure weights are on cpu
+
+        _synchronize_device(self.device)
+        _clean_memory_on_device(self.device)
+
+    def wait_for_block(self, block_idx: int):
+        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
+            return
+        self._wait_blocks_move(block_idx)
+
+    def submit_move_blocks_forward(self, blocks: list[nn.Module], block_idx: int):
+        # check if blocks_to_swap is enabled
+        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
+            return
+
+        if not self.forward_only:
+            # if backward is enabled, we do not swap blocks in forward pass more than blocks_to_swap, because it should be on GPU
+            if block_idx >= self.blocks_to_swap:
+                return
+            block_idx_to_cpu = block_idx
+            block_idx_to_cuda = self.num_blocks - self.blocks_to_swap + block_idx
+            block_idx_to_cuda = block_idx_to_cuda % self.num_blocks  # this does nothing for backward offloading
+            self._submit_move_blocks(blocks, block_idx_to_cpu, block_idx_to_cuda)
+            return
+
+        # We use two strategies here for forward-only offloading:
+        # 1. If blocks_to_swap is less than half of num_blocks, we swap the num_blocks blocks without wrapping around.
+        #   This reduces the number of swaps, so it is especially useful for small blocks_to_swap or lightweight models like Qwen-Image
+        # 2. If blocks_to_swap is more than half of num_blocks, we swap the blocks with wrapping around.
+        #   This is the common strategy used in most offloading implementations. It transfers all blocks in a wrapping manner.
+        #   This is useful for large blocks_to_swap or heavyweight models like Wan/HunyuanVideo, where the transfer time is less significant compared to computation time.
+
+        # current block to swap out (to CPU)
+        block_idx_to_cpu = block_idx
+
+        if self.blocks_to_swap < (self.num_blocks // 2):
+            # strategy 1: no wrap around
+            # If the current block is in the middle blocks that are not swapped, do nothing
+            if self.blocks_to_swap <= block_idx < self.num_blocks - self.blocks_to_swap:
+                return
+            if block_idx < self.blocks_to_swap:
+                # move the next block to cuda
+                block_idx_to_cuda = (self.num_blocks - self.blocks_to_swap + block_idx) % self.num_blocks
+            else:
+                # move the previous block to cuda
+                block_idx_to_cuda = block_idx - (self.num_blocks - self.blocks_to_swap)
+        else:
+            # strategy 2: with wrap around
+            block_idx_to_cuda = self.num_blocks - self.blocks_to_swap + block_idx
+            block_idx_to_cuda = block_idx_to_cuda % self.num_blocks  # this works for forward-only offloading
+
+        self._submit_move_blocks(blocks, block_idx_to_cpu, block_idx_to_cuda)
+
+
+class _DirectCopier:
+    """
+    Transfer engine for ``LoRAStreamOffloader`` when the masters are *pinned*: copies a flat pinned host
+    buffer straight into a flat device buffer on a private copy stream (async Host->Device only).
+
+    This is the seam that isolates all CUDA stream/event primitives from the residency logic. The residency
+    side only ever calls ``submit`` / ``wait``; ``_StagedCopier`` (pageable master + worker thread) and a
+    future non-CUDA device shim plug in behind the same interface without the residency code changing. Keep
+    every ``torch.cuda.*`` transfer primitive inside copiers like this one.
+    """
+
+    def __init__(self, device: torch.device, debug: bool = False):
+        self.device = device
+        self.debug = debug
+        self.copy_stream = torch.cuda.Stream(device=device)
+        self._events = {}  # key -> cuda.Event recording H2D completion
+        self._xfer = {}  # key -> (ev_a, ev_b) transfer-timing events (debug only)
+
+    def submit(self, key, dst_flat: torch.Tensor, src_flat: torch.Tensor, gate_event=None):
+        """Begin the H2D copy of ``src_flat`` into ``dst_flat``, gated behind ``gate_event`` (compute done with dst)."""
+        ev_a = ev_b = None
+        with torch.cuda.stream(self.copy_stream):
+            if gate_event is not None:
+                self.copy_stream.wait_event(gate_event)  # do not overwrite a buffer still used by compute
+            if self.debug:
+                ev_a = torch.cuda.Event(enable_timing=True)
+                ev_b = torch.cuda.Event(enable_timing=True)
+                ev_a.record(self.copy_stream)
+            dst_flat.copy_(src_flat, non_blocking=True)  # src is pinned -> async H2D
+            if self.debug:
+                ev_b.record(self.copy_stream)
+                self._xfer[key] = (ev_a, ev_b)
+        self._events[key] = self.copy_stream.record_event()
+
+    def wait(self, key):
+        """Return the H2D-completion event for ``key`` (None if nothing is in flight for it)."""
+        return self._events.get(key)
+
+    def pop_xfer_timing(self, key):
+        """Pop the (start, end) transfer-timing events recorded for ``key`` (debug only)."""
+        return self._xfer.pop(key, None)
+
+    def sync(self):
+        self.copy_stream.synchronize()
+
+    def reset(self):
+        self._events.clear()
+        self._xfer.clear()
+
+
+class _StagedCopier:
+    """
+    Transfer engine for ``LoRAStreamOffloader`` when the masters are kept in *pageable* host memory: each
+    H2D is staged through a small pool of pinned buffers filled by a background worker thread. Same
+    ``submit`` / ``wait`` interface as ``_DirectCopier``, so the residency code is identical.
+
+    Why: pinning every streaming master (what ``_DirectCopier`` needs) can cost tens of GB of non-pageable
+    "shared GPU memory" (notably on Windows). Here the masters stay pageable and only ``num_staging`` pinned
+    buffers exist. A pageable->device copy is, inside the CUDA driver, already a synchronous serial
+    pageable->pinned->device transfer; doing it ourselves on a worker thread lets the pageable->pinned
+    memcpy overlap the previous block's pinned->device DMA and -- crucially -- keeps it off the main thread
+    so compute kernel launches are never blocked. That is the whole speed difference versus the old
+    synchronous pageable path this replaces.
+
+    Hazards handled:
+      - a staging buffer is reused only after the H2D that consumed it has completed (one consume-event per
+        buffer; the worker waits on it before overwriting);
+      - ``wait`` blocks the caller until the worker has actually enqueued the H2D and recorded its
+        completion event, so the residency layer can ``current_stream().wait_event`` it as usual.
+    """
+
+    def __init__(self, device: torch.device, num_staging: int = 2, debug: bool = False):
+        self.device = device
+        self.num_staging = max(1, num_staging)
+        self.debug = debug
+        self.copy_stream = torch.cuda.Stream(device=device)
+        self._pool = ThreadPoolExecutor(max_workers=1)
+        self._futures = {}  # key -> Future -> (load_event, (ev_a, ev_b))
+        self._xfer = {}  # key -> (ev_a, ev_b) transfer-timing events (debug only)
+        self._staging = None  # list[pinned uint8 tensor] (lazily sized to the flat block)
+        self._staging_free = None  # list[cuda.Event | None]: the H2D that last consumed each buffer
+        self._rr = 0  # round-robin staging index (touched only on the worker thread)
+
+    def _ensure_staging(self, nbytes: int):
+        if self._staging is None:
+            self._staging = [torch.empty(nbytes, dtype=torch.uint8).pin_memory(device=self.device) for _ in range(self.num_staging)]
+            self._staging_free = [None] * self.num_staging
+
+    def _run(self, dst_flat: torch.Tensor, src_flat: torch.Tensor, gate_event):
+        # runs on the worker thread: pageable->pinned memcpy, then enqueue the async pinned->device H2D
+        idx = self._rr
+        self._rr = (idx + 1) % self.num_staging
+        staging = self._staging[idx]
+
+        free_ev = self._staging_free[idx]
+        if free_ev is not None:
+            free_ev.synchronize()  # CPU(worker) waits until the prior H2D released this buffer -- off main thread
+
+        staging.copy_(src_flat)  # pageable -> pinned, plain CPU memcpy
+
+        ev_a = ev_b = None
+        with torch.cuda.stream(self.copy_stream):
+            if gate_event is not None:
+                self.copy_stream.wait_event(gate_event)  # do not overwrite a dst still used by compute
+            if self.debug:
+                ev_a = torch.cuda.Event(enable_timing=True)
+                ev_b = torch.cuda.Event(enable_timing=True)
+                ev_a.record(self.copy_stream)
+            dst_flat.copy_(staging, non_blocking=True)  # pinned -> device, async
+            if self.debug:
+                ev_b.record(self.copy_stream)
+            done = self.copy_stream.record_event()  # H2D complete == this staging buffer is free again
+        self._staging_free[idx] = done
+        return done, (ev_a, ev_b)
+
+    def submit(self, key, dst_flat: torch.Tensor, src_flat: torch.Tensor, gate_event=None):
+        self._ensure_staging(src_flat.numel())  # flat uint8 -> numel == byte size
+        self._futures[key] = self._pool.submit(self._run, dst_flat, src_flat, gate_event)
+
+    def wait(self, key):
+        fut = self._futures.get(key)
+        if fut is None:
+            return None
+        load_event, xfer = fut.result()  # block until the worker enqueued the H2D (tiny if it ran ahead)
+        if self.debug and xfer[0] is not None:
+            self._xfer[key] = xfer
+        return load_event
+
+    def pop_xfer_timing(self, key):
+        return self._xfer.pop(key, None)
+
+    def _flush(self):
+        for fut in list(self._futures.values()):
+            fut.result()
+
+    def sync(self):
+        self._flush()
+        self.copy_stream.synchronize()
+
+    def reset(self):
+        self._flush()
+        self.copy_stream.synchronize()  # ensure all in-flight H2D are done before clearing staging state
+        self._futures.clear()
+        self._xfer.clear()
+        if self._staging is not None:
+            self._staging_free = [None] * self.num_staging
+
+    def __del__(self):
+        pool = getattr(self, "_pool", None)
+        if pool is not None:
+            pool.shutdown(wait=False)
+
+
+class LoRAStreamOffloader:
+    """
+    H2D-only offloader for training where the base weights are frozen (e.g. LoRA / LoHa / LoKr).
+
+    The classic block swap (``ModelOffloader``) *exchanges* a block: it copies one block back to the
+    CPU (Device->Host) while copying the next one in (Host->Device). When the base weights never change,
+    the CPU already holds an identical copy, so the D2H half is pure overhead. This offloader keeps a
+    permanent (pinned) master copy of every streamed weight on the CPU and only ever copies Host->Device,
+    removing the D2H transfer and the CPU-side staging memcpy entirely.
+
+    Layout (N = num_blocks, S = blocks_to_swap, B = ring_size):
+      - ``N - S`` blocks stay resident on the GPU permanently.
+      - ``S`` "streaming" blocks are placed at even intervals; their ``Linear`` weights live on the CPU
+        (pinned) as masters and are streamed into a small ring of ``B`` GPU buffers (default 2: one being
+        computed on, one being prefetched). Non-Linear params / buffers of streaming blocks stay resident.
+
+    Mechanics:
+      - Because the streamed weights are frozen (no ``.grad`` to preserve), we swap the whole ``module.weight``
+        *reference* between its CPU master Parameter and a preallocated GPU ring Parameter, instead of the
+        ``.data`` storage trick used by ``ModelOffloader``. No per-step ``cudaMalloc``.
+      - All swap weights of a block are packed into a single contiguous (byte) buffer -- one pinned CPU
+        buffer per streaming block and one GPU buffer per ring slot -- and the individual weights are
+        dtype/shape views into it. Loading a block is therefore a single H2D ``copy_`` instead of one
+        per Linear, which removes the per-tensor launch/driver overhead on PCIe/CPU-bound systems.
+      - Eviction is free: repoint ``module.weight`` back to the CPU master (no transfer). A ring slot is
+        overwritten only when its block is consumed, so the last ``B`` streamed blocks of a pass stay
+        resident. On the reverse pass they are exactly the first blocks needed, so the forward<->backward
+        turn-around needs no reload.
+      - Direction is implicit in the callback: ``submit_move_blocks_forward`` prefetches the block ``B``
+        streaming-slots *ahead* (larger index); the backward hook prefetches ``B`` slots *behind* (smaller
+        index). No explicit scheduler / no ``prepare_*_before_backward``.
+      - ``wait_for_block`` self-heals: if a needed streaming block is not resident in its slot (cold start,
+        or an inference pass boundary), it is loaded on demand.
+    """
+
+    def __init__(
+        self,
+        block_type: str,
+        blocks: list[nn.Module],
+        num_blocks: int,
+        blocks_to_swap: int,
+        supports_backward: bool,
+        device: torch.device,
+        ring_size: int = 2,
+        use_pinned_memory: bool = True,
+        debug: bool = False,
+    ):
+        self.block_type = block_type
+        self._blocks = blocks
+        self.num_blocks = num_blocks  # N
+        self.blocks_to_swap = blocks_to_swap  # requested S
+        self.device = device
+        self.use_pinned_memory = use_pinned_memory
+        self.supports_backward = supports_backward
+        self.forward_only = not supports_backward
+
+        import os
+
+        if not debug:
+            debug = os.getenv("MUSUBI_TUNER_OFFLOADER_DEBUG", "0") == "1"
+        self.debug = debug
+        self.debug_interval = int(os.getenv("MUSUBI_TUNER_OFFLOADER_DEBUG_INTERVAL", "10"))  # steps between prints
+
+        assert device.type == "cuda", "LoRAStreamOffloader currently supports CUDA only"
+
+        # ---- streaming placement: S evenly spaced block indices (midpoint formula -> distinct for S <= N) ----
+        stream_idx = sorted({((2 * i + 1) * num_blocks) // (2 * blocks_to_swap) for i in range(blocks_to_swap)})
+        self.stream_idx = stream_idx
+        self.S = len(stream_idx)  # actual streaming count (>=1; dedup is a no-op unless S is close to N)
+        self.rank = {b: k for k, b in enumerate(stream_idx)}  # block_idx -> position in stream_idx
+        self.is_stream = [b in self.rank for b in range(num_blocks)]
+        self.B = min(ring_size, self.S)  # clamp ring to number of streaming blocks
+
+        # ---- finetuning guard: H2D-only must never lose weight updates ----
+        for b in stream_idx:
+            for m in self._swap_modules(blocks[b]):
+                assert not m.weight.requires_grad, (
+                    "LoRAStreamOffloader requires frozen base weights (LoRA / no full fine-tune). "
+                    f"Found a trainable Linear weight in block {b}."
+                )
+
+        # ---- transfer engine: owns the copy stream and all H2D primitives ----
+        # pinned masters -> direct async H2D; pageable masters -> stage through a worker thread + pinned pool
+        # (low "shared GPU memory" footprint, e.g. on Windows). Both expose the same submit/wait interface.
+        if use_pinned_memory:
+            self.copier = _DirectCopier(device, debug=self.debug)
+        else:
+            self.copier = _StagedCopier(device, num_staging=self.B, debug=self.debug)
+
+        # ---- runtime state (GPU buffers allocated lazily in prepare_block_devices_before_forward) ----
+        self.cpu_master = {}  # block_idx -> [CPU (pinned) Parameter per swap weight] (views into cpu_flat)
+        self.cpu_flat = {}  # block_idx -> flat (pinned) uint8 CPU tensor backing the masters
+        self.ring_param = None  # [slot] -> [GPU nn.Parameter per swap weight] (views into ring_flat)
+        self.ring_flat = None  # [slot] -> flat uint8 GPU tensor backing the ring params
+        self._layout = None  # ([byte offset per swap weight], total bytes) shared by all streaming blocks
+        self.in_slot = [None] * self.B  # slot -> block_idx currently bound to this slot (or None)
+        self.free_event = [None] * self.B  # slot -> cuda.Event recording when compute finished using the slot
+        self._module_cache = {}  # block_idx -> [modules with swap weights]
+
+        self._wait_ctx = "fwd"  # context tag for wait/load timing ("fwd" while in the forward loop, "bwd" in hooks)
+        if self.debug:
+            self._dbg_reset()
+
+        # ---- backward hooks: only where there is work (prefetch this block, or wait the previous one) ----
+        if supports_backward:
+            self.remove_handles = []
+            for i, block in enumerate(blocks):
+                hook = self._create_backward_hook(i)
+                if hook is not None:
+                    self.remove_handles.append(block.register_full_backward_hook(hook))
+
+        print(
+            f"LoRAStreamOffloader[{block_type}]: H2D-only block swap. "
+            f"{self.S} streaming / {num_blocks} blocks, ring={self.B}, pinned={use_pinned_memory}. "
+            f"streaming indices: {stream_idx}"
+        )
+
+    # ------------------------------------------------------------------ helpers
+
+    def _swap_modules(self, block: nn.Module) -> list[nn.Module]:
+        # same selection rule as ModelOffloader.swap_weight_devices_cuda: Linear layers with a weight
+        mods = []
+        for _, m in block.named_modules():
+            if hasattr(m, "weight") and m.weight is not None and m.__class__.__name__.endswith("Linear"):
+                mods.append(m)
+        return mods
+
+    def _modules(self, block_idx: int) -> list[nn.Module]:
+        cached = self._module_cache.get(block_idx)
+        if cached is None:
+            cached = self._swap_modules(self._blocks[block_idx])
+            self._module_cache[block_idx] = cached
+        return cached
+
+    def _bind(self, block_idx: int, params: list[nn.Parameter]):
+        for m, p in zip(self._modules(block_idx), params):
+            m.weight = p
+
+    @staticmethod
+    def _compute_layout(weights: list[torch.Tensor]) -> tuple[list[int], int]:
+        """Byte offsets (aligned so ``view(dtype)`` is valid for any element size) and total size of the flat buffer."""
+        align = 256
+        offsets = []
+        total = 0
+        for w in weights:
+            total = (total + align - 1) // align * align
+            offsets.append(total)
+            total += w.numel() * w.element_size()
+        return offsets, total
+
+    def _flat_views(self, flat: torch.Tensor, weights: list[torch.Tensor]) -> list[torch.Tensor]:
+        """dtype/shape views into a flat uint8 buffer, one per swap weight, following ``self._layout``."""
+        offsets, _ = self._layout
+        return [flat[off : off + w.numel() * w.element_size()].view(w.dtype).view(w.shape) for off, w in zip(offsets, weights)]
+
+    def _load(self, rank: int, slot: int, ctx: str = "fwd"):
+        """Stream the streaming block at ``rank`` into ring ``slot`` (H2D only), repointing weights."""
+        blk = self.stream_idx[rank]
+
+        # already resident (e.g. ring survived a re-prepare): the slot content is still identical to the
+        # immutable CPU master, so just rebind the weights to the ring -- no transfer needed
+        if self.in_slot[slot] == blk:
+            self._bind(blk, self.ring_param[slot])
+            return
+
+        # evict the slot's current owner: repoint its weights back to the CPU master (no transfer)
+        prev = self.in_slot[slot]
+        if prev is not None:
+            self._bind(prev, self.cpu_master[prev])
+
+        # single coalesced H2D copy of the whole block, gated behind the slot's free_event (compute done
+        # with it). The ring params are views into ring_flat, so the copy bumps their autograd version
+        # counter: safe under gradient checkpointing (weights are re-read at recompute time, and a slot is
+        # only overwritten after its free_event); without checkpointing autograd raises a version error
+        # instead of silently using stale-saved weights -- which the old per-tensor `.data.copy_` hid.
+        self.copier.submit(blk, self.ring_flat[slot], self.cpu_flat[blk], self.free_event[slot])
+
+        self._bind(blk, self.ring_param[slot])  # reference swap (no .data trick, no cudaMalloc)
+        self.in_slot[slot] = blk
+
+        if self.debug:
+            c = self._cur[ctx]
+            c["loads"] += 1
+            xfer = self.copier.pop_xfer_timing(blk)
+            if xfer is not None:
+                c["xfer_ev"].append(xfer)
+
+    # ------------------------------------------------------------------ public interface (mirrors ModelOffloader)
+
+    def set_forward_only(self, forward_only: bool):
+        self.copier.sync()
+        self.forward_only = forward_only
+
+    def __del__(self):
+        if getattr(self, "supports_backward", False):
+            for handle in getattr(self, "remove_handles", []):
+                handle.remove()
+
+    def prepare_block_devices_before_forward(self, blocks: list[nn.Module]):
+        if self.S == 0:
+            return
+
+        if self.debug:
+            print(f"[{self.block_type}] Prepare block devices before forward (H2D-only)")
+
+        first_time = not self.cpu_master  # device placement is only needed on the very first call
+        cpu_device = torch.device("cpu")
+        for i, b in enumerate(blocks):
+            if not self.is_stream[i]:
+                if first_time:
+                    b.to(self.device)  # resident block: buffers + Linear weights on GPU permanently
+                    weighs_to_device(b, self.device)
+                continue
+
+            if first_time:
+                # move the whole block to device (buffers, norms, bias), then pull the swap weights back to
+                # the CPU into one flat (pinned) buffer per block as the persistent masters
+                b.to(self.device)
+                mods = self._modules(i)
+                weights = [m.weight.data for m in mods]
+                if self._layout is None:
+                    self._layout = self._compute_layout(weights)  # first streaming block defines the shared layout
+                flat = torch.empty(self._layout[1], dtype=torch.uint8, device=cpu_device)
+                if self.use_pinned_memory:
+                    flat = flat.pin_memory(device=self.device)
+                master = []
+                for m, view in zip(mods, self._flat_views(flat, weights)):
+                    view.copy_(m.weight.data)  # one-time D2H into the flat master
+                    m.weight.data = view
+                    master.append(m.weight)  # keep the original Parameter object as the persistent master
+                self.cpu_flat[i] = flat
+                self.cpu_master[i] = master
+            else:
+                # re-prepare (e.g. around in-training sampling): the non-swap parts are already on the device
+                # and the CPU masters have stayed on the CPU throughout, so just repoint weights to the masters.
+                # IMPORTANT: do NOT call b.to(device) here -- it would drag the CPU masters onto the GPU and
+                # they would never be released (every swap block would end up resident).
+                self._bind(i, self.cpu_master[i])
+
+        # validate homogeneous streaming blocks (shared ring template)
+        template = self.cpu_master[self.stream_idx[0]]
+        for b in self.stream_idx[1:]:
+            assert len(self.cpu_master[b]) == len(template), f"block {b} has a different number of swap weights"
+            for p, t in zip(self.cpu_master[b], template):
+                assert p.data.shape == t.data.shape and p.data.dtype == t.data.dtype, (
+                    f"block {b} swap-weight shape/dtype differs from the streaming template"
+                )
+
+        # preallocate the GPU ring (once); copies happen into these flat buffers, never reallocated
+        if self.ring_param is None:
+            template_weights = [p.data for p in template]
+            self.ring_flat = [torch.empty(self._layout[1], dtype=torch.uint8, device=self.device) for _ in range(self.B)]
+            self.ring_param = [
+                [nn.Parameter(view, requires_grad=False) for view in self._flat_views(flat, template_weights)]
+                for flat in self.ring_flat
+            ]
+
+        # reset transfer state and preload the first B streaming blocks. in_slot deliberately survives:
+        # after a completed backward pass the ring already holds ranks 0..B-1, so these loads become
+        # no-op rebinds (see _load) instead of redundant cold-start transfers.
+        self.free_event = [None] * self.B
+        self.copier.reset()
+        for k in range(self.B):
+            self._load(k, k)
+
+        _synchronize_device(self.device)
+        _clean_memory_on_device(self.device)
+
+    def wait_for_block(self, block_idx: int):
+        if self.S == 0 or not self.is_stream[block_idx]:
+            return
+        ctx = self._wait_ctx
+        if self.debug and ctx == "fwd" and block_idx == self.stream_idx[0]:
+            self._dbg_boundary()  # a forward pass is starting -> finalize the previous step's stats
+        j = self.rank[block_idx]
+        slot = j % self.B
+        if self.in_slot[slot] != block_idx:
+            # self-healing: not resident (cold start / inference pass boundary / same-direction restart)
+            if self.debug:
+                self._cur[ctx]["self_heal"] += 1
+            self._load(j, slot, ctx)
+        ev = self.copier.wait(block_idx)
+        if ev is None:
+            return
+        if not self.debug:
+            torch.cuda.current_stream().wait_event(ev)  # compute waits for H2D (GPU-side, does not block CPU)
+            return
+        # debug: measure the actual GPU stall (idle time the compute stream spends on wait_event)
+        s = torch.cuda.current_stream()
+        c = self._cur[ctx]
+        c["waits"] += 1
+        if not ev.query():
+            c["not_ready"] += 1  # H2D not finished when compute reached it -> a real stall (cheap, no sync)
+        ev_a = torch.cuda.Event(enable_timing=True)
+        ev_b = torch.cuda.Event(enable_timing=True)
+        ev_a.record(s)
+        s.wait_event(ev)
+        ev_b.record(s)
+        c["stall_ev"].append((ev_a, ev_b))
+        # staged copier records its transfer-timing events on the worker thread, so they only become
+        # available once wait() has joined it; the direct copier already popped them in _load (returns None)
+        xfer = self.copier.pop_xfer_timing(block_idx)
+        if xfer is not None:
+            c["xfer_ev"].append(xfer)
+
+    def submit_move_blocks_forward(self, blocks: list[nn.Module], block_idx: int):
+        # consumed block_idx in the forward pass: free its slot and prefetch the block B streaming-slots ahead
+        if self.S == 0 or not self.is_stream[block_idx]:
+            return
+        j = self.rank[block_idx]
+        self.free_event[j % self.B] = torch.cuda.current_stream().record_event()
+        if j + self.B < self.S:
+            self._load(j + self.B, (j + self.B) % self.B, "fwd")  # same slot as rank j; evicts rank j
+        elif self.forward_only and j == self.S - 1 and self.S > self.B:
+            # inference: no backward pass will consume the retained tail, and the next forward pass starts
+            # from rank 0 again. Wrap around: all tail slots are free once the last streaming block is
+            # consumed, so preload the first B ranks now. The H2D overlaps with the remaining (resident)
+            # blocks and the output layers, instead of stalling at the start of the next pass (self-heal).
+            for k in range(self.B):
+                self._load(k, k, "fwd")
+
+    def _create_backward_hook(self, block_index: int):
+        prefetch = self.is_stream[block_index]
+        wait_prev = block_index - 1 >= 0 and self.is_stream[block_index - 1]
+        if not prefetch and not wait_prev:
+            return None
+
+        def backward_hook(module, grad_input, grad_output):
+            if prefetch:
+                # consumed block_index in backward: free its slot and prefetch B streaming-slots behind
+                j = self.rank[block_index]
+                self.free_event[j % self.B] = torch.cuda.current_stream().record_event()
+                if j - self.B >= 0:
+                    self._load(j - self.B, (j - self.B) % self.B, "bwd")  # same slot as rank j; evicts rank j
+            if wait_prev:
+                self._wait_ctx = "bwd"
+                self.wait_for_block(block_index - 1)  # ensure the next block to run is ready
+                self._wait_ctx = "fwd"
+            return None
+
+        return backward_hook
+
+    # ------------------------------------------------------------------ debug timing
+
+    def _dbg_reset(self):
+        self._dbg_empty = lambda: {"stall_ev": [], "xfer_ev": [], "waits": 0, "not_ready": 0, "self_heal": 0, "loads": 0}
+        self._cur = {"fwd": self._dbg_empty(), "bwd": self._dbg_empty()}
+        self._roll = {"fwd": defaultdict(float), "bwd": defaultdict(float)}
+        self._step = 0
+
+    def _dbg_boundary(self):
+        # finalize the events accumulated for the step that just ended, then print every debug_interval steps
+        if not any(self._cur[c]["waits"] or self._cur[c]["loads"] for c in ("fwd", "bwd")):
+            return  # very first forward: nothing finished yet
+        torch.cuda.synchronize()  # ensure timing events are complete (debug only)
+        for ctx in ("fwd", "bwd"):
+            c, r = self._cur[ctx], self._roll[ctx]
+            r["stall_ms"] += sum(a.elapsed_time(b) for a, b in c["stall_ev"])
+            r["xfer_ms"] += sum(a.elapsed_time(b) for a, b in c["xfer_ev"])
+            for k in ("waits", "not_ready", "self_heal", "loads"):
+                r[k] += c[k]
+            self._cur[ctx] = self._dbg_empty()
+        self._step += 1
+        if self._step % self.debug_interval == 0:
+            self._dbg_print()
+
+    def _dbg_print(self):
+        n = self.debug_interval
+        print(f"[{self.block_type}] H2D-only timing (avg over {n} steps):")
+        for ctx in ("fwd", "bwd"):
+            r = self._roll[ctx]
+            print(
+                f"  {ctx}: stall {r['stall_ms'] / n:6.2f} ms/step, H2D {r['xfer_ms'] / n:6.2f} ms/step"
+                f" | waits {r['waits'] / n:.1f}, not-ready {r['not_ready'] / n:.1f},"
+                f" self-heal {r['self_heal'] / n:.1f}, loads {r['loads'] / n:.1f}"
+            )
+            self._roll[ctx] = defaultdict(float)
diff --git a/musubi_tuner/modules/fp8_optimization_utils.py b/musubi_tuner/modules/fp8_optimization_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..264ffc6f080554e7c8c258e9eb4101b2a0770b55
--- /dev/null
+++ b/musubi_tuner/modules/fp8_optimization_utils.py
@@ -0,0 +1,564 @@
+import os
+from typing import List, Optional, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import logging
+
+from tqdm import tqdm
+
+from musubi_tuner.utils.safetensors_utils import MemoryEfficientSafeOpen, TensorWeightAdapter, WeightTransformHooks
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+from musubi_tuner.utils.device_utils import clean_memory_on_device
+
+
+def calculate_fp8_maxval(exp_bits=4, mantissa_bits=3, sign_bits=1):
+    """
+    Calculate the maximum representable value in FP8 format.
+    Default is E4M3 format (4-bit exponent, 3-bit mantissa, 1-bit sign). Only supports E4M3 and E5M2 with sign bit.
+
+    Args:
+        exp_bits (int): Number of exponent bits
+        mantissa_bits (int): Number of mantissa bits
+        sign_bits (int): Number of sign bits (0 or 1)
+
+    Returns:
+        float: Maximum value representable in FP8 format
+    """
+    assert exp_bits + mantissa_bits + sign_bits == 8, "Total bits must be 8"
+    if exp_bits == 4 and mantissa_bits == 3 and sign_bits == 1:
+        return torch.finfo(torch.float8_e4m3fn).max
+    elif exp_bits == 5 and mantissa_bits == 2 and sign_bits == 1:
+        return torch.finfo(torch.float8_e5m2).max
+    else:
+        raise ValueError(f"Unsupported FP8 format: E{exp_bits}M{mantissa_bits} with sign_bits={sign_bits}")
+
+
+# The following is a manual calculation method (wrong implementation for E5M2), kept for reference.
+"""
+# Calculate exponent bias
+bias = 2 ** (exp_bits - 1) - 1
+
+# Calculate maximum mantissa value
+mantissa_max = 1.0
+for i in range(mantissa_bits - 1):
+    mantissa_max += 2 ** -(i + 1)
+
+# Calculate maximum value
+max_value = mantissa_max * (2 ** (2**exp_bits - 1 - bias))
+
+return max_value
+"""
+
+
+def quantize_fp8(tensor, scale, fp8_dtype, max_value, min_value):
+    """
+    Quantize a tensor to FP8 format using PyTorch's native FP8 dtype support.
+
+    Args:
+        tensor (torch.Tensor): Tensor to quantize
+        scale (float or torch.Tensor): Scale factor
+        fp8_dtype (torch.dtype): Target FP8 dtype (torch.float8_e4m3fn or torch.float8_e5m2)
+        max_value (float): Maximum representable value in FP8
+        min_value (float): Minimum representable value in FP8
+
+    Returns:
+        torch.Tensor: Quantized tensor in FP8 format
+    """
+    tensor = tensor.to(torch.float32)  # ensure tensor is in float32 for division
+
+    # Create scaled tensor
+    tensor = torch.div(tensor, scale).nan_to_num_(0.0)  # handle NaN values, equivalent to nonzero_mask in previous function
+
+    # Clamp tensor to range
+    tensor = tensor.clamp_(min=min_value, max=max_value)
+
+    # Convert to FP8 dtype
+    tensor = tensor.to(fp8_dtype)
+
+    return tensor
+
+
+def optimize_state_dict_with_fp8(
+    state_dict: dict,
+    calc_device: Union[str, torch.device],
+    target_layer_keys: Optional[list[str]] = None,
+    exclude_layer_keys: Optional[list[str]] = None,
+    exp_bits: int = 4,
+    mantissa_bits: int = 3,
+    move_to_device: bool = False,
+    quantization_mode: str = "block",
+    block_size: Optional[int] = 64,
+):
+    """
+    Optimize Linear layer weights in a model's state dict to FP8 format. The state dict is modified in-place.
+    This function is a static version of load_safetensors_with_fp8_optimization without loading from files.
+
+    Args:
+        state_dict (dict): State dict to optimize, replaced in-place
+        calc_device (str): Device to quantize tensors on
+        target_layer_keys (list, optional): Layer key patterns to target (None for all Linear layers)
+        exclude_layer_keys (list, optional): Layer key patterns to exclude
+        exp_bits (int): Number of exponent bits
+        mantissa_bits (int): Number of mantissa bits
+        move_to_device (bool): Move optimized tensors to the calculating device
+
+    Returns:
+        dict: FP8 optimized state dict
+    """
+    if exp_bits == 4 and mantissa_bits == 3:
+        fp8_dtype = torch.float8_e4m3fn
+    elif exp_bits == 5 and mantissa_bits == 2:
+        fp8_dtype = torch.float8_e5m2
+    else:
+        raise ValueError(f"Unsupported FP8 format: E{exp_bits}M{mantissa_bits}")
+
+    # Calculate FP8 max value
+    max_value = calculate_fp8_maxval(exp_bits, mantissa_bits)
+    min_value = -max_value  # this function supports only signed FP8
+
+    # Create optimized state dict
+    optimized_count = 0
+
+    # Enumerate tarket keys
+    target_state_dict_keys = []
+    for key in state_dict.keys():
+        # Check if it's a weight key and matches target patterns
+        is_target = (target_layer_keys is None or any(pattern in key for pattern in target_layer_keys)) and key.endswith(".weight")
+        is_excluded = exclude_layer_keys is not None and any(pattern in key for pattern in exclude_layer_keys)
+        is_target = is_target and not is_excluded
+
+        if is_target and isinstance(state_dict[key], torch.Tensor):
+            target_state_dict_keys.append(key)
+
+    # Process each key
+    for key in tqdm(target_state_dict_keys):
+        value = state_dict[key]
+
+        # Save original device and dtype
+        original_device = value.device
+        original_dtype = value.dtype
+
+        # Move to calculation device
+        if calc_device is not None:
+            value = value.to(calc_device)
+
+        quantized_weight, scale_tensor = quantize_weight(key, value, fp8_dtype, max_value, min_value, quantization_mode, block_size)
+
+        # Add to state dict using original key for weight and new key for scale
+        fp8_key = key  # Maintain original key
+        scale_key = key.replace(".weight", ".scale_weight")
+
+        if not move_to_device:
+            quantized_weight = quantized_weight.to(original_device)
+
+        # keep scale shape: [1] or [out,1] or [out, num_blocks, 1]. We can determine the quantization mode from the shape of scale_weight in the patched model.
+        scale_tensor = scale_tensor.to(dtype=original_dtype, device=quantized_weight.device)
+
+        state_dict[fp8_key] = quantized_weight
+        state_dict[scale_key] = scale_tensor
+
+        optimized_count += 1
+
+        if calc_device is not None:  # optimized_count % 10 == 0 and
+            # free memory on calculation device
+            clean_memory_on_device(calc_device)
+
+    logger.info(f"Number of optimized Linear layers: {optimized_count}")
+    return state_dict
+
+
+def quantize_weight(
+    key: str,
+    tensor: torch.Tensor,
+    fp8_dtype: torch.dtype,
+    max_value: float,
+    min_value: float,
+    quantization_mode: str = "block",
+    block_size: int = 64,
+):
+    original_shape = tensor.shape
+
+    # Determine quantization mode
+    if quantization_mode == "block":
+        if tensor.ndim != 2:
+            quantization_mode = "tensor"  # fallback to per-tensor
+        else:
+            out_features, in_features = tensor.shape
+            if in_features % block_size != 0:
+                quantization_mode = "channel"  # fallback to per-channel
+                logger.warning(
+                    f"Layer {key} with shape {tensor.shape} is not divisible by block_size {block_size}, fallback to per-channel quantization."
+                )
+            else:
+                num_blocks = in_features // block_size
+                tensor = tensor.contiguous().view(out_features, num_blocks, block_size)  # [out, num_blocks, block_size]
+    elif quantization_mode == "channel":
+        if tensor.ndim != 2:
+            quantization_mode = "tensor"  # fallback to per-tensor
+
+    # Calculate scale factor (per-tensor or per-output-channel with percentile or max)
+    # value shape is expected to be [out_features, in_features] for Linear weights
+    if quantization_mode == "channel" or quantization_mode == "block":
+        # row-wise percentile to avoid being dominated by outliers
+        # result shape: [out_features, 1] or [out_features, num_blocks, 1]
+        scale_dim = 1 if quantization_mode == "channel" else 2
+        abs_w = torch.abs(tensor)
+
+        # shape: [out_features, 1] or [out_features, num_blocks, 1]
+        row_max = torch.max(abs_w, dim=scale_dim, keepdim=True).values
+        scale = row_max / max_value
+
+    else:
+        # per-tensor
+        tensor_max = torch.max(torch.abs(tensor).view(-1))
+        scale = tensor_max / max_value
+
+    # print(f"Optimizing {key} with scale: {scale}")
+
+    # numerical safety
+    scale = torch.clamp(scale, min=1e-8)
+    scale = scale.to(torch.float32)  # ensure scale is in float32 for division
+
+    # Quantize weight to FP8 (scale can be scalar or [out,1], broadcasting works)
+    quantized_weight = quantize_fp8(tensor, scale, fp8_dtype, max_value, min_value)
+
+    # If block-wise, restore original shape
+    if quantization_mode == "block":
+        quantized_weight = quantized_weight.view(original_shape)  # restore to original shape [out, in]
+
+    return quantized_weight, scale
+
+
+def load_safetensors_with_fp8_optimization(
+    model_files: List[str],
+    calc_device: Union[str, torch.device],
+    target_layer_keys=None,
+    exclude_layer_keys=None,
+    exp_bits=4,
+    mantissa_bits=3,
+    move_to_device=False,
+    weight_hook=None,
+    quantization_mode: str = "block",
+    block_size: Optional[int] = 64,
+    disable_numpy_memmap: bool = False,
+    weight_transform_hooks: Optional[WeightTransformHooks] = None,
+    allow_prequantized_fp8: bool = False,
+) -> dict:
+    """
+    Load weight tensors from safetensors files and merge LoRA weights into the state dict with explicit FP8 optimization.
+
+    Args:
+        model_files (list[str]): List of model files to load
+        calc_device (str or torch.device): Device to quantize tensors on
+        target_layer_keys (list, optional): Layer key patterns to target for optimization (None for all Linear layers)
+        exclude_layer_keys (list, optional): Layer key patterns to exclude from optimization
+        exp_bits (int): Number of exponent bits
+        mantissa_bits (int): Number of mantissa bits
+        move_to_device (bool): Move optimized tensors to the calculating device
+        weight_hook (callable, optional): Function to apply to each weight tensor before optimization
+        quantization_mode (str): Quantization mode, "tensor", "channel", or "block"
+        block_size (int, optional): Block size for block-wise quantization (used if quantization_mode is "block")
+        disable_numpy_memmap (bool): Disable numpy memmap when loading safetensors
+        weight_transform_hooks (WeightTransformHooks, optional): Hooks for weight transformation during loading
+        allow_prequantized_fp8 (bool): If True, target weights that are already FP8 are kept as-is instead of
+            raising; their scale is expected to be supplied separately (e.g. via weight_transform_hooks).
+
+    Returns:
+        dict: FP8 optimized state dict
+    """
+    if exp_bits == 4 and mantissa_bits == 3:
+        fp8_dtype = torch.float8_e4m3fn
+    elif exp_bits == 5 and mantissa_bits == 2:
+        fp8_dtype = torch.float8_e5m2
+    else:
+        raise ValueError(f"Unsupported FP8 format: E{exp_bits}M{mantissa_bits}")
+
+    # Calculate FP8 max value
+    max_value = calculate_fp8_maxval(exp_bits, mantissa_bits)
+    min_value = -max_value  # this function supports only signed FP8
+
+    # Define function to determine if a key is a target key. target means fp8 optimization, not for weight hook.
+    def is_target_key(key):
+        # Check if weight key matches target patterns and does not match exclude patterns
+        is_target = (target_layer_keys is None or any(pattern in key for pattern in target_layer_keys)) and key.endswith(".weight")
+        is_excluded = exclude_layer_keys is not None and any(pattern in key for pattern in exclude_layer_keys)
+        return is_target and not is_excluded
+
+    # Create optimized state dict
+    optimized_count = 0
+
+    # Process each file
+    state_dict = {}
+    for model_file in model_files:
+        with MemoryEfficientSafeOpen(model_file, disable_numpy_memmap=disable_numpy_memmap) as original_f:
+            f = TensorWeightAdapter(weight_transform_hooks, original_f) if weight_transform_hooks is not None else original_f
+
+            keys = f.keys()
+            for key in tqdm(keys, desc=f"Loading {os.path.basename(model_file)}", unit="key"):
+                value = f.get_tensor(key)
+
+                # Save original device
+                original_device = value.device  # usually cpu
+
+                if weight_hook is not None:
+                    # Apply weight hook if provided
+                    value = weight_hook(key, value, keep_on_calc_device=(calc_device is not None))
+
+                if not is_target_key(key):
+                    target_device = calc_device if (calc_device is not None and move_to_device) else original_device
+                    value = value.to(target_device)
+                    state_dict[key] = value
+                    continue
+
+                # Move to calculation device
+                if calc_device is not None:
+                    value = value.to(calc_device)
+
+                original_dtype = value.dtype
+                if original_dtype.itemsize == 1:
+                    if allow_prequantized_fp8 and original_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+                        # Weight is already FP8-quantized (e.g. a pre-quantized ComfyUI checkpoint).
+                        # Keep it as-is without re-quantizing; the matching scale is supplied separately
+                        # (e.g. renamed to `.scale_weight` via weight_transform_hooks).
+                        if not move_to_device:
+                            value = value.to(original_device)
+                        state_dict[key] = value
+                        optimized_count += 1
+                        continue
+                    raise ValueError(
+                        f"Layer {key} is already in {original_dtype} format. `--fp8_scaled` optimization should not be applied. Please use fp16/bf16/float32 model weights."
+                        + f" / レイヤー {key} は既に{original_dtype}形式です。`--fp8_scaled` 最適化は適用できません。FP16/BF16/Float32のモデル重みを使用してください。"
+                    )
+                quantized_weight, scale_tensor = quantize_weight(
+                    key, value, fp8_dtype, max_value, min_value, quantization_mode, block_size
+                )
+
+                # Add to state dict using original key for weight and new key for scale
+                fp8_key = key  # Maintain original key
+                scale_key = key.replace(".weight", ".scale_weight")
+                assert fp8_key != scale_key, "FP8 key and scale key must be different"
+
+                if not move_to_device:
+                    quantized_weight = quantized_weight.to(original_device)
+
+                # keep scale shape: [1] or [out,1] or [out, num_blocks, 1]. We can determine the quantization mode from the shape of scale_weight in the patched model.
+                scale_tensor = scale_tensor.to(dtype=original_dtype, device=quantized_weight.device)
+
+                state_dict[fp8_key] = quantized_weight
+                state_dict[scale_key] = scale_tensor
+
+                optimized_count += 1
+
+                if calc_device is not None and optimized_count % 10 == 0:
+                    # free memory on calculation device
+                    clean_memory_on_device(calc_device)
+
+    logger.info(f"Number of optimized Linear layers: {optimized_count}")
+    return state_dict
+
+
+def fp8_linear_forward_patch(self: nn.Linear, x, use_scaled_mm=False, max_value=None):
+    """
+    Patched forward method for Linear layers with FP8 weights.
+
+    Args:
+        self: Linear layer instance
+        x (torch.Tensor): Input tensor
+        use_scaled_mm (bool): Use scaled_mm for FP8 Linear layers, requires SM 8.9+ (RTX 40 series)
+        max_value (float): Maximum value for FP8 quantization. If None, no quantization is applied for input tensor.
+
+    Returns:
+        torch.Tensor: Result of linear transformation
+    """
+    if use_scaled_mm:
+        # **not tested**
+        # _scaled_mm only works for per-tensor scale for now (per-channel scale does not work in certain cases)
+        if self.scale_weight.ndim != 1:
+            raise ValueError("scaled_mm only supports per-tensor scale_weight for now.")
+
+        input_dtype = x.dtype
+        original_weight_dtype = self.scale_weight.dtype
+        target_dtype = self.weight.dtype
+        # assert x.ndim == 3, "Input tensor must be 3D (batch_size, seq_len, hidden_dim)"
+
+        if max_value is None:
+            # no input quantization
+            scale_x = torch.tensor(1.0, dtype=torch.float32, device=x.device)
+        else:
+            # calculate scale factor for input tensor
+            scale_x = (torch.max(torch.abs(x.flatten())) / max_value).to(torch.float32)
+
+            # quantize input tensor to FP8: this seems to consume a lot of memory
+            fp8_max_value = torch.finfo(target_dtype).max
+            fp8_min_value = torch.finfo(target_dtype).min
+            x = quantize_fp8(x, scale_x, target_dtype, fp8_max_value, fp8_min_value)
+
+        original_shape = x.shape
+        x = x.reshape(-1, x.shape[-1]).to(target_dtype)
+
+        weight = self.weight.t()
+        scale_weight = self.scale_weight.to(torch.float32)
+
+        if self.bias is not None:
+            # float32 is not supported with bias in scaled_mm
+            o = torch._scaled_mm(x, weight, out_dtype=original_weight_dtype, bias=self.bias, scale_a=scale_x, scale_b=scale_weight)
+        else:
+            o = torch._scaled_mm(x, weight, out_dtype=input_dtype, scale_a=scale_x, scale_b=scale_weight)
+
+        o = o.reshape(original_shape[0], original_shape[1], -1) if len(original_shape) == 3 else o.reshape(original_shape[0], -1)
+        return o.to(input_dtype)
+
+    else:
+        # Dequantize the weight
+        original_dtype = self.scale_weight.dtype
+        if self.scale_weight.ndim < 3:
+            # per-tensor or per-channel quantization, we can broadcast
+            dequantized_weight = self.weight.to(original_dtype) * self.scale_weight
+        else:
+            # block-wise quantization, need to reshape weight to match scale shape for broadcasting
+            out_features, num_blocks, _ = self.scale_weight.shape
+            dequantized_weight = self.weight.to(original_dtype).contiguous().view(out_features, num_blocks, -1)
+            dequantized_weight = dequantized_weight * self.scale_weight
+            dequantized_weight = dequantized_weight.view(self.weight.shape)
+
+        # Perform linear transformation
+        if self.bias is not None:
+            output = F.linear(x, dequantized_weight, self.bias)
+        else:
+            output = F.linear(x, dequantized_weight)
+
+        return output
+
+
+def apply_fp8_monkey_patch(model, optimized_state_dict, use_scaled_mm=False):
+    """
+    Apply monkey patching to a model using FP8 optimized state dict.
+
+    Args:
+        model (nn.Module): Model instance to patch
+        optimized_state_dict (dict): FP8 optimized state dict
+        use_scaled_mm (bool): Use scaled_mm for FP8 Linear layers, requires SM 8.9+ (RTX 40 series)
+
+    Returns:
+        nn.Module: The patched model (same instance, modified in-place)
+    """
+    # # Calculate FP8 float8_e5m2 max value
+    # max_value = calculate_fp8_maxval(5, 2)
+    max_value = None  # do not quantize input tensor
+
+    # Find all scale keys to identify FP8-optimized layers
+    scale_keys = [k for k in optimized_state_dict.keys() if k.endswith(".scale_weight")]
+
+    # Enumerate patched layers
+    patched_module_paths = set()
+    scale_shape_info = {}
+    for scale_key in scale_keys:
+        # Extract module path from scale key (remove .scale_weight)
+        module_path = scale_key.rsplit(".scale_weight", 1)[0]
+        patched_module_paths.add(module_path)
+
+        # Store scale shape information
+        scale_shape_info[module_path] = optimized_state_dict[scale_key].shape
+
+    patched_count = 0
+
+    # Apply monkey patch to each layer with FP8 weights
+    for name, module in model.named_modules():
+        # Check if this module has a corresponding scale_weight
+        has_scale = name in patched_module_paths
+
+        # Apply patch if it's a Linear layer with FP8 scale
+        if isinstance(module, nn.Linear) and has_scale:
+            # register the scale_weight as a buffer to load the state_dict
+            # module.register_buffer("scale_weight", torch.tensor(1.0, dtype=module.weight.dtype))
+            scale_shape = scale_shape_info[name]
+            module.register_buffer("scale_weight", torch.ones(scale_shape, dtype=module.weight.dtype))
+
+            # Create a new forward method with the patched version.
+            def new_forward(self, x):
+                return fp8_linear_forward_patch(self, x, use_scaled_mm, max_value)
+
+            # Bind method to module
+            module.forward = new_forward.__get__(module, type(module))
+
+            patched_count += 1
+
+    logger.info(f"Number of monkey-patched Linear layers: {patched_count}")
+    return model
+
+
+# Example usage
+def example_usage():
+    # Small test model
+    class TestModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            fc1 = nn.Linear(768, 3072)
+            act1 = nn.GELU()
+            fc2 = nn.Linear(3072, 768)
+            act2 = nn.GELU()
+            fc3 = nn.Linear(768, 768)
+
+            # Set layer names for testing
+            self.single_blocks = nn.ModuleList([fc1, act1, fc2, act2, fc3])
+
+            self.fc4 = nn.Linear(768, 128)
+
+        def forward(self, x):
+            for layer in self.single_blocks:
+                x = layer(x)
+            x = self.fc4(x)
+            return x
+
+    # Instantiate model
+    test_model = TestModel()
+    test_model.to(torch.float16)  # convert to FP16 for testing
+
+    # Test input tensor
+    test_input = torch.randn(1, 768, dtype=torch.float16)
+
+    # Calculate output before optimization
+    with torch.no_grad():
+        original_output = test_model(test_input)
+        print("original output", original_output[0, :5])
+
+    # Get state dict
+    state_dict = test_model.state_dict()
+
+    # Apply FP8 optimization to state dict
+    cuda_device = torch.device("cuda")
+    optimized_state_dict = optimize_state_dict_with_fp8(state_dict, cuda_device, ["single_blocks"], ["2"])
+
+    # Apply monkey patching to the model
+    optimized_model = TestModel()  # re-instantiate model
+    optimized_model.to(torch.float16)  # convert to FP16 for testing
+    apply_fp8_monkey_patch(optimized_model, optimized_state_dict)
+
+    # Load optimized state dict
+    optimized_model.load_state_dict(optimized_state_dict, strict=True, assign=True)  # assign=True to load buffer
+
+    # Calculate output after optimization
+    with torch.no_grad():
+        optimized_output = optimized_model(test_input)
+        print("optimized output", optimized_output[0, :5])
+
+    # Compare accuracy
+    error = torch.mean(torch.abs(original_output - optimized_output))
+    print(f"Mean absolute error: {error.item()}")
+
+    # Check memory usage
+    original_params = sum(p.nelement() * p.element_size() for p in test_model.parameters()) / (1024 * 1024)
+    print(f"Model parameter memory: {original_params:.2f} MB")
+    optimized_params = sum(p.nelement() * p.element_size() for p in optimized_model.parameters()) / (1024 * 1024)
+    print(f"Optimized model parameter memory: {optimized_params:.2f} MB")
+
+    return test_model
+
+
+if __name__ == "__main__":
+    example_usage()
diff --git a/musubi_tuner/modules/lr_schedulers.py b/musubi_tuner/modules/lr_schedulers.py
new file mode 100644
index 0000000000000000000000000000000000000000..61edb04f9dfebcf8e2930a377350aa3742d19e79
--- /dev/null
+++ b/musubi_tuner/modules/lr_schedulers.py
@@ -0,0 +1,67 @@
+import torch
+
+
+class RexLR(torch.optim.lr_scheduler.LRScheduler):
+    """
+    Reflected Exponential (REX) learning rate scheduler (https://arxiv.org/abs/2107.04197)
+    Modified from: https://github.com/IvanVassi/REX_LR (Apache-2.0 License)
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer to schedule the learning rate for
+        max_lr (float): The maximum learning rate
+        min_lr (float): The minimum learning rate
+        num_steps (int): The total number of training steps
+        num_warmup_steps (int): The number of warmup steps
+        rex_alpha (float): Constant added to the denominator of the REX factor;
+            prevents division-by-zero and softens the initial decay (default: 0.1).
+        rex_beta (float): Multiplier of z in the denominator of the REX factor;
+            controls how quickly the decay flattens as z increases (default: 0.9).
+        last_epoch (int): The index of the last step
+    """
+
+    def __init__(self, optimizer, max_lr, min_lr=0.0, num_steps=0, num_warmup_steps=0, rex_alpha=0.1, rex_beta=0.9, last_epoch=-1):
+        if min_lr > max_lr:
+            raise ValueError(f'Value of "min_lr" should be less than value of "max_lr". Got min_lr={min_lr} and max_lr={max_lr}')
+        if num_warmup_steps > num_steps:
+            raise ValueError(f"num_warmup_steps ({num_warmup_steps}) must be less than or equal to num_steps ({num_steps})")
+
+        self.min_lr = min_lr
+        self.max_lr = max_lr
+        self.num_steps = num_steps
+        self.num_warmup_steps = num_warmup_steps
+        self.rex_alpha = rex_alpha
+        self.rex_beta = rex_beta
+        self.last_epoch = last_epoch
+
+        # Ensure each parameter group has an "initial_lr" key to avoid issues when resuming
+        for group in optimizer.param_groups:
+            group.setdefault("initial_lr", group["lr"])
+
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        # Single warmup step
+        if self.num_warmup_steps == 1 and self.last_epoch == 1:
+            return [self.min_lr for _ in self.base_lrs]
+        # Multiple warmup steps; increase lr linearly from min_lr to max_lr
+        elif self.num_warmup_steps > 1 and self.last_epoch >= 1 and self.last_epoch <= (self.num_warmup_steps - 1):
+            return [
+                self.min_lr + (self.max_lr - self.min_lr) * (self.last_epoch - 1) / (self.num_warmup_steps - 1)
+                for _ in self.base_lrs
+            ]
+
+        # Post-warmup phase: adjust step relative to the end of warmup
+        step_after = self.last_epoch - self.num_warmup_steps
+        remaining_steps = self.num_steps - self.num_warmup_steps
+
+        # Avoid LR spiking
+        if step_after >= remaining_steps or step_after == -1 or remaining_steps <= 0:
+            return [self.min_lr for _ in self.base_lrs]
+
+        # Calculate REX curve for current step
+        rex_z = (remaining_steps - (step_after % remaining_steps)) / remaining_steps
+        rex_factor = self.min_lr / self.max_lr + (1.0 - self.min_lr / self.max_lr) * (
+            rex_z / (self.rex_alpha + self.rex_beta * rex_z)
+        )
+
+        return [base_lr * rex_factor for base_lr in self.base_lrs]
diff --git a/musubi_tuner/modules/scheduling_flow_match_discrete.py b/musubi_tuner/modules/scheduling_flow_match_discrete.py
new file mode 100644
index 0000000000000000000000000000000000000000..c507ec4eb050463188e250c20aec8d1fde2c4a5d
--- /dev/null
+++ b/musubi_tuner/modules/scheduling_flow_match_discrete.py
@@ -0,0 +1,257 @@
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        reverse (`bool`, defaults to `True`):
+            Whether to reverse the timestep schedule.
+    """
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        reverse: bool = True,
+        solver: str = "euler",
+        n_tokens: Optional[int] = None,
+    ):
+        sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
+
+        if not reverse:
+            sigmas = sigmas.flip(0)
+
+        self.sigmas = sigmas
+        # the value fed to model
+        self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.supported_solver = ["euler"]
+        if solver not in self.supported_solver:
+            raise ValueError(
+                f"Solver {solver} not supported. Supported solvers: {self.supported_solver}"
+            )
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Union[str, torch.device] = None,
+        n_tokens: int = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            n_tokens (`int`, *optional*):
+                Number of tokens in the input sequence.
+        """
+        self.num_inference_steps = num_inference_steps
+        
+        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
+        sigmas = self.sd3_time_shift(sigmas)
+
+        if not self.config.reverse:
+            sigmas = 1 - sigmas
+
+        self.sigmas = sigmas
+        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
+            dtype=torch.float32, device=device
+        )
+
+        # Reset step index
+        self._step_index = None
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: Optional[int] = None
+    ) -> torch.Tensor:
+        return sample
+
+    def sd3_time_shift(self, t: torch.Tensor):
+        return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        return_dict: bool = True,
+    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            n_tokens (`int`, *optional*):
+                Number of tokens in the input sequence.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+                tuple.
+
+        Returns:
+            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+                returned, otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if (
+            isinstance(timestep, int)
+            or isinstance(timestep, torch.IntTensor)
+            or isinstance(timestep, torch.LongTensor)
+        ):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
+
+        if self.config.solver == "euler":
+            prev_sample = sample + model_output.to(torch.float32) * dt
+        else:
+            raise ValueError(
+                f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}"
+            )
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self):
+        return self.config.num_train_timesteps
diff --git a/musubi_tuner/modules/unet_causal_3d_blocks.py b/musubi_tuner/modules/unet_causal_3d_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..27d544170ece6a370cdacfe9e31367b884c2e516
--- /dev/null
+++ b/musubi_tuner/modules/unet_causal_3d_blocks.py
@@ -0,0 +1,818 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+
+from diffusers.utils import logging
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import SpatialNorm
+from diffusers.models.attention_processor import Attention
+from diffusers.models.normalization import AdaGroupNorm
+from diffusers.models.normalization import RMSNorm
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None):
+    seq_len = n_frame * n_hw
+    mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
+    for i in range(seq_len):
+        i_frame = i // n_hw
+        mask[i, : (i_frame + 1) * n_hw] = 0
+    if batch_size is not None:
+        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+    return mask
+
+
+class CausalConv3d(nn.Module):
+    """
+    Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations.
+    This maintains temporal causality in video generation tasks.
+    """
+
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        stride: Union[int, Tuple[int, int, int]] = 1,
+        dilation: Union[int, Tuple[int, int, int]] = 1,
+        pad_mode="replicate",
+        chunk_size=0,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.pad_mode = pad_mode
+        padding = (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size - 1, 0)  # W, H, T
+        self.time_causal_padding = padding
+        self.chunk_size = chunk_size
+
+        self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+    def original_forward(self, x):
+        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+        return self.conv(x)
+
+    def forward(self, x):
+        if self.chunk_size == 0:
+            return self.original_forward(x)
+
+        # if not large, call original forward
+        if x.shape[4] < self.chunk_size * 1.5:
+            return self.original_forward(x)
+
+        # # debug: verify the original forward is the same as chunked forward
+        # orig_forwarded_value = None
+        # if x.shape[4] < self.chunk_size * 4:
+        #     orig_forwarded_value = self.original_forward(x)
+
+        # get the kernel size
+        kernel_size = self.conv.kernel_size[0]  # assume cubic kernel
+        assert kernel_size == self.conv.kernel_size[1] == self.conv.kernel_size[2], "Only cubic kernels are supported"
+        padding_size = kernel_size // 2  # 1 for kernel_size=3, 0 for kernel_size=1
+
+        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+
+        B, C, D, H, W = orig_shape = x.shape
+        chunk_size = self.chunk_size
+        chunk_size -= chunk_size % self.conv.stride[2]  # make sure the chunk size is divisible by stride
+        # print(f"chunked forward: {x.shape}, chunk_size: {chunk_size}")
+
+        # calculate the indices for chunking with overlap and padding by kernel size and stride
+        indices = []
+        i = 0
+        while i < W - padding_size:
+            start_idx = i - padding_size
+            end_idx = min(i + chunk_size + padding_size, W)
+            if i == 0:
+                start_idx = 0
+                end_idx += padding_size  # to make sure the first chunk is divisible by stride
+            if W - end_idx < chunk_size // 2:  # small chunk at the end
+                end_idx = W
+            indices.append((start_idx, end_idx))
+            i = end_idx - padding_size
+        # print(f"chunked forward: {x.shape}, chunked indices: {indices}")
+
+        chunks = []
+        for start_idx, end_idx in indices:
+            chunk = x[:, :, :, :, start_idx:end_idx]
+            chunk_output = self.conv(chunk)
+            # print(chunk.shape, chunk_output.shape)
+            chunks.append(chunk_output)
+
+        # concatenate the chunks
+        x = torch.cat(chunks, dim=4)
+
+        assert (
+            x.shape[2] == ((D - padding_size * 2) + self.conv.stride[0] - 1) // self.conv.stride[0]
+        ), f"Invalid shape: {x.shape}, {orig_shape}, {padding_size}, {self.conv.stride}"
+        assert (
+            x.shape[3] == ((H - padding_size * 2) + self.conv.stride[1] - 1) // self.conv.stride[1]
+        ), f"Invalid shape: {x.shape}, {orig_shape}, {padding_size}, {self.conv.stride}"
+        assert (
+            x.shape[4] == ((W - padding_size * 2) + self.conv.stride[2] - 1) // self.conv.stride[2]
+        ), f"Invalid shape: {x.shape}, {orig_shape}, {padding_size}, {self.conv.stride}"
+
+        # # debug: verify the original forward is the same as chunked forward
+        # if orig_forwarded_value is not None:
+        #     assert torch.allclose(
+        #         orig_forwarded_value, x, rtol=1e-4, atol=1e-2
+        #     ), f"Chunked forward is different from original forward. {x.shape}, {orig_shape}, {padding_size}, {self.conv.stride}, {self.conv.kernel_size}"
+
+        return x
+
+
+class UpsampleCausal3D(nn.Module):
+    """
+    A 3D upsampling layer with an optional convolution.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool = False,
+        use_conv_transpose: bool = False,
+        out_channels: Optional[int] = None,
+        name: str = "conv",
+        kernel_size: Optional[int] = None,
+        padding=1,
+        norm_type=None,
+        eps=None,
+        elementwise_affine=None,
+        bias=True,
+        interpolate=True,
+        upsample_factor=(2, 2, 2),
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+        self.interpolate = interpolate
+        self.upsample_factor = upsample_factor
+
+        if norm_type == "ln_norm":
+            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(channels, eps, elementwise_affine)
+        elif norm_type is None:
+            self.norm = None
+        else:
+            raise ValueError(f"unknown norm_type: {norm_type}")
+
+        conv = None
+        if use_conv_transpose:
+            raise NotImplementedError
+        elif use_conv:
+            if kernel_size is None:
+                kernel_size = 3
+            conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias)
+
+        if name == "conv":
+            self.conv = conv
+        else:
+            self.Conv2d_0 = conv
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        output_size: Optional[int] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.norm is not None:
+            raise NotImplementedError
+
+        if self.use_conv_transpose:
+            return self.conv(hidden_states)
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            hidden_states = hidden_states.contiguous()
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if self.interpolate:
+            B, C, T, H, W = hidden_states.shape
+            first_h, other_h = hidden_states.split((1, T - 1), dim=2)
+            if output_size is None:
+                if T > 1:
+                    other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest")
+
+                first_h = first_h.squeeze(2)
+                first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest")
+                first_h = first_h.unsqueeze(2)
+            else:
+                raise NotImplementedError
+
+            if T > 1:
+                hidden_states = torch.cat((first_h, other_h), dim=2)
+            else:
+                hidden_states = first_h
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        if self.use_conv:
+            if self.name == "conv":
+                hidden_states = self.conv(hidden_states)
+            else:
+                hidden_states = self.Conv2d_0(hidden_states)
+
+        return hidden_states
+
+
+class DownsampleCausal3D(nn.Module):
+    """
+    A 3D downsampling layer with an optional convolution.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool = False,
+        out_channels: Optional[int] = None,
+        padding: int = 1,
+        name: str = "conv",
+        kernel_size=3,
+        norm_type=None,
+        eps=None,
+        elementwise_affine=None,
+        bias=True,
+        stride=2,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = stride
+        self.name = name
+
+        if norm_type == "ln_norm":
+            self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(channels, eps, elementwise_affine)
+        elif norm_type is None:
+            self.norm = None
+        else:
+            raise ValueError(f"unknown norm_type: {norm_type}")
+
+        if use_conv:
+            conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, bias=bias)
+        else:
+            raise NotImplementedError
+
+        if name == "conv":
+            self.Conv2d_0 = conv
+            self.conv = conv
+        elif name == "Conv2d_0":
+            self.conv = conv
+        else:
+            self.conv = conv
+
+    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+        assert hidden_states.shape[1] == self.channels
+
+        if self.norm is not None:
+            hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+        assert hidden_states.shape[1] == self.channels
+
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class ResnetBlockCausal3D(nn.Module):
+    r"""
+    A Resnet block.
+    """
+
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        groups: int = 32,
+        groups_out: Optional[int] = None,
+        pre_norm: bool = True,
+        eps: float = 1e-6,
+        non_linearity: str = "swish",
+        skip_time_act: bool = False,
+        # default, scale_shift, ada_group, spatial
+        time_embedding_norm: str = "default",
+        kernel: Optional[torch.FloatTensor] = None,
+        output_scale_factor: float = 1.0,
+        use_in_shortcut: Optional[bool] = None,
+        up: bool = False,
+        down: bool = False,
+        conv_shortcut_bias: bool = True,
+        conv_3d_out_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.up = up
+        self.down = down
+        self.output_scale_factor = output_scale_factor
+        self.time_embedding_norm = time_embedding_norm
+        self.skip_time_act = skip_time_act
+
+        linear_cls = nn.Linear
+
+        if groups_out is None:
+            groups_out = groups
+
+        if self.time_embedding_norm == "ada_group":
+            self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)
+        elif self.time_embedding_norm == "spatial":
+            self.norm1 = SpatialNorm(in_channels, temb_channels)
+        else:
+            self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+        self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1)
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == "default":
+                self.time_emb_proj = linear_cls(temb_channels, out_channels)
+            elif self.time_embedding_norm == "scale_shift":
+                self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels)
+            elif self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+                self.time_emb_proj = None
+            else:
+                raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ")
+        else:
+            self.time_emb_proj = None
+
+        if self.time_embedding_norm == "ada_group":
+            self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)
+        elif self.time_embedding_norm == "spatial":
+            self.norm2 = SpatialNorm(out_channels, temb_channels)
+        else:
+            self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+
+        self.dropout = torch.nn.Dropout(dropout)
+        conv_3d_out_channels = conv_3d_out_channels or out_channels
+        self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1)
+
+        self.nonlinearity = get_activation(non_linearity)
+
+        self.upsample = self.downsample = None
+        if self.up:
+            self.upsample = UpsampleCausal3D(in_channels, use_conv=False)
+        elif self.down:
+            self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op")
+
+        self.use_in_shortcut = self.in_channels != conv_3d_out_channels if use_in_shortcut is None else use_in_shortcut
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = CausalConv3d(
+                in_channels,
+                conv_3d_out_channels,
+                kernel_size=1,
+                stride=1,
+                bias=conv_shortcut_bias,
+            )
+
+    def forward(
+        self,
+        input_tensor: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        hidden_states = input_tensor
+
+        if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+            hidden_states = self.norm1(hidden_states, temb)
+        else:
+            hidden_states = self.norm1(hidden_states)
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:
+            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+            if hidden_states.shape[0] >= 64:
+                input_tensor = input_tensor.contiguous()
+                hidden_states = hidden_states.contiguous()
+            input_tensor = self.upsample(input_tensor, scale=scale)
+            hidden_states = self.upsample(hidden_states, scale=scale)
+        elif self.downsample is not None:
+            input_tensor = self.downsample(input_tensor, scale=scale)
+            hidden_states = self.downsample(hidden_states, scale=scale)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if self.time_emb_proj is not None:
+            if not self.skip_time_act:
+                temb = self.nonlinearity(temb)
+            temb = self.time_emb_proj(temb, scale)[:, :, None, None]
+
+        if temb is not None and self.time_embedding_norm == "default":
+            hidden_states = hidden_states + temb
+
+        if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+            hidden_states = self.norm2(hidden_states, temb)
+        else:
+            hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == "scale_shift":
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+def get_down_block3d(
+    down_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    temb_channels: int,
+    add_downsample: bool,
+    downsample_stride: int,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    downsample_padding: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    downsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+):
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
+    if down_block_type == "DownEncoderBlockCausal3D":
+        return DownEncoderBlockCausal3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            dropout=dropout,
+            add_downsample=add_downsample,
+            downsample_stride=downsample_stride,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block3d(
+    up_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    prev_output_channel: int,
+    temb_channels: int,
+    add_upsample: bool,
+    upsample_scale_factor: Tuple,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    resolution_idx: Optional[int] = None,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    upsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+) -> nn.Module:
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+    if up_block_type == "UpDecoderBlockCausal3D":
+        return UpDecoderBlockCausal3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            resolution_idx=resolution_idx,
+            dropout=dropout,
+            add_upsample=add_upsample,
+            upsample_scale_factor=upsample_scale_factor,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            temb_channels=temb_channels,
+        )
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlockCausal3D(nn.Module):
+    """
+    A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        attn_groups: Optional[int] = None,
+        resnet_pre_norm: bool = True,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        self.add_attention = add_attention
+
+        if attn_groups is None:
+            attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlockCausal3D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        attentions = []
+
+        if attention_head_dim is None:
+            logger.warn(
+                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
+            )
+            attention_head_dim = in_channels
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    Attention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        rescale_output_factor=output_scale_factor,
+                        eps=resnet_eps,
+                        norm_num_groups=attn_groups,
+                        spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
+                        residual_connection=True,
+                        bias=True,
+                        upcast_softmax=True,
+                        _from_deprecated_attn_block=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+    def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                B, C, T, H, W = hidden_states.shape
+                hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c")
+                attention_mask = prepare_causal_attention_mask(T, H * W, hidden_states.dtype, hidden_states.device, batch_size=B)
+                hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask)
+                hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W)
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class DownEncoderBlockCausal3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor: float = 1.0,
+        add_downsample: bool = True,
+        downsample_stride: int = 2,
+        downsample_padding: int = 1,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=None,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    DownsampleCausal3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                        stride=downsample_stride,
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=None, scale=scale)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states, scale)
+
+        return hidden_states
+
+
+class UpDecoderBlockCausal3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        resolution_idx: Optional[int] = None,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor: float = 1.0,
+        add_upsample: bool = True,
+        upsample_scale_factor=(2, 2, 2),
+        temb_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlockCausal3D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [
+                    UpsampleCausal3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        upsample_factor=upsample_scale_factor,
+                    )
+                ]
+            )
+        else:
+            self.upsamplers = None
+
+        self.resolution_idx = resolution_idx
+
+    def forward(
+        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
+    ) -> torch.FloatTensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
diff --git a/musubi_tuner/utils/__init__.py b/musubi_tuner/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/musubi_tuner/utils/device_utils.py b/musubi_tuner/utils/device_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..479808846adafa8230906000e974b863aa5b731c
--- /dev/null
+++ b/musubi_tuner/utils/device_utils.py
@@ -0,0 +1,28 @@
+from typing import Optional, Union
+import torch
+
+
+def clean_memory_on_device(device: Optional[Union[str, torch.device]]):
+    if device is None:
+        return
+    if isinstance(device, str):
+        device = torch.device(device)
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    elif device.type == "cpu":
+        pass
+    elif device.type == "mps":  # not tested
+        torch.mps.empty_cache()
+
+
+def synchronize_device(device: Optional[Union[str, torch.device]]):
+    if device is None:
+        return
+    if isinstance(device, str):
+        device = torch.device(device)
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    elif device.type == "xpu":
+        torch.xpu.synchronize()
+    elif device.type == "mps":
+        torch.mps.synchronize()
diff --git a/musubi_tuner/utils/huggingface_utils.py b/musubi_tuner/utils/huggingface_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..0dc7bd7dbb2ef70e0b6244b9db686aae00f46408
--- /dev/null
+++ b/musubi_tuner/utils/huggingface_utils.py
@@ -0,0 +1,89 @@
+import threading
+from typing import Union, BinaryIO
+from huggingface_hub import HfApi
+from pathlib import Path
+import argparse
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+def fire_in_thread(f, *args, **kwargs):
+    threading.Thread(target=f, args=args, kwargs=kwargs).start()
+
+
+def exists_repo(repo_id: str, repo_type: str, revision: str = "main", token: str = None):
+    api = HfApi(
+        token=token,
+    )
+    try:
+        api.repo_info(repo_id=repo_id, revision=revision, repo_type=repo_type)
+        return True
+    except:
+        return False
+
+
+def upload(
+    args: argparse.Namespace,
+    src: Union[str, Path, bytes, BinaryIO],
+    dest_suffix: str = "",
+    force_sync_upload: bool = False,
+):
+    repo_id = args.huggingface_repo_id
+    repo_type = args.huggingface_repo_type
+    token = args.huggingface_token
+    path_in_repo = args.huggingface_path_in_repo + dest_suffix if args.huggingface_path_in_repo is not None else None
+    private = args.huggingface_repo_visibility is None or args.huggingface_repo_visibility != "public"
+    api = HfApi(token=token)
+    if not exists_repo(repo_id=repo_id, repo_type=repo_type, token=token):
+        try:
+            api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private)
+        except Exception as e:  # RepositoryNotFoundError or something else
+            logger.error("===========================================")
+            logger.error(f"failed to create HuggingFace repo / HuggingFaceのリポジトリの作成に失敗しました : {e}")
+            logger.error("===========================================")
+
+    is_folder = (type(src) == str and os.path.isdir(src)) or (isinstance(src, Path) and src.is_dir())
+
+    def uploader():
+        try:
+            if is_folder:
+                api.upload_folder(
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    folder_path=src,
+                    path_in_repo=path_in_repo,
+                )
+            else:
+                api.upload_file(
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    path_or_fileobj=src,
+                    path_in_repo=path_in_repo,
+                )
+        except Exception as e:  # RuntimeError or something else
+            logger.error("===========================================")
+            logger.error(f"failed to upload to HuggingFace / HuggingFaceへのアップロードに失敗しました : {e}")
+            logger.error("===========================================")
+
+    if args.async_upload and not force_sync_upload:
+        fire_in_thread(uploader)
+    else:
+        uploader()
+
+
+def list_dir(
+    repo_id: str,
+    subfolder: str,
+    repo_type: str,
+    revision: str = "main",
+    token: str = None,
+):
+    api = HfApi(
+        token=token,
+    )
+    repo_info = api.repo_info(repo_id=repo_id, revision=revision, repo_type=repo_type)
+    file_list = [file for file in repo_info.siblings if file.rfilename.startswith(subfolder)]
+    return file_list
diff --git a/musubi_tuner/utils/image_utils.py b/musubi_tuner/utils/image_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec872b3b4cbdf7392dc65c33bf4de83ca0b01fa7
--- /dev/null
+++ b/musubi_tuner/utils/image_utils.py
@@ -0,0 +1,40 @@
+import numpy as np
+import torch
+from PIL import Image
+from typing import Tuple, Optional
+
+from musubi_tuner.dataset import image_video_dataset
+
+
+# prepare image
+def preprocess_image(
+    image: Image, w: int, h: int, handle_alpha: bool = False
+) -> Tuple[torch.Tensor, np.ndarray, Optional[np.ndarray]]:
+    """
+    Preprocess the image for the model.
+    Args:
+        image (Image): The input image. RGB or RGBA format.
+        w (int): The target bucket width.
+        h (int): The target bucket height.
+        handle_alpha (bool): Whether to handle alpha channel for tensor and numpy array.
+    Returns:
+        Tuple[torch.Tensor, np.ndarray, Optional[np.ndarray]]:
+            - image_tensor: The preprocessed image tensor (NCHW format). -1.0 to 1.0.
+            - image_np: The original image as a numpy array (HWC format). 0 to 255.
+            - alpha: The alpha channel of the image if present in original size, otherwise None.
+    """
+    if image.mode == "RGBA":
+        alpha = image.split()[-1]
+    else:
+        alpha = None
+    if handle_alpha:
+        image = image.convert("RGBA")
+    else:
+        image = image.convert("RGB")
+
+    image_np = np.array(image)  # PIL to numpy, HWC
+
+    image_np = image_video_dataset.resize_image_to_bucket(image_np, (w, h))  # TODO move this to this file
+    image_tensor = torch.from_numpy(image_np).float() / 127.5 - 1.0  # -1 to 1.0, HWC
+    image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0)  # HWC -> CHW -> NCHW, N=1
+    return image_tensor, image_np, alpha
diff --git a/musubi_tuner/utils/lora_utils.py b/musubi_tuner/utils/lora_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3a005a3f7bd7d08790ae73dbd0f78c410c38cad
--- /dev/null
+++ b/musubi_tuner/utils/lora_utils.py
@@ -0,0 +1,300 @@
+import os
+import re
+from typing import Dict, List, Optional, Union
+import torch
+
+import logging
+
+from tqdm import tqdm
+
+from musubi_tuner.utils.device_utils import synchronize_device
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+from musubi_tuner.modules.fp8_optimization_utils import load_safetensors_with_fp8_optimization
+from musubi_tuner.utils.safetensors_utils import (
+    MemoryEfficientSafeOpen,
+    TensorWeightAdapter,
+    WeightTransformHooks,
+    get_split_weight_filenames,
+)
+
+
+def detect_network_type(lora_sd: Dict[str, torch.Tensor]) -> str:
+    """Detect network type (lora, loha, lokr) from state dict keys."""
+    for key in lora_sd:
+        if "lora_down" in key:
+            return "lora"
+        if "hada_w1_a" in key:
+            return "loha"
+        if "lokr_w1" in key:
+            return "lokr"
+    return "lora"  # default
+
+
+def filter_lora_state_dict(
+    weights_sd: Dict[str, torch.Tensor],
+    include_pattern: Optional[str] = None,
+    exclude_pattern: Optional[str] = None,
+) -> Dict[str, torch.Tensor]:
+    # apply include/exclude patterns
+    original_key_count = len(weights_sd.keys())
+    if include_pattern is not None:
+        regex_include = re.compile(include_pattern)
+        weights_sd = {k: v for k, v in weights_sd.items() if regex_include.search(k)}
+        logger.info(f"Filtered keys with include pattern {include_pattern}: {original_key_count} -> {len(weights_sd.keys())}")
+
+    if exclude_pattern is not None:
+        original_key_count_ex = len(weights_sd.keys())
+        regex_exclude = re.compile(exclude_pattern)
+        weights_sd = {k: v for k, v in weights_sd.items() if not regex_exclude.search(k)}
+        logger.info(f"Filtered keys with exclude pattern {exclude_pattern}: {original_key_count_ex} -> {len(weights_sd.keys())}")
+
+    if len(weights_sd) != original_key_count:
+        remaining_keys = list(set([k.split(".", 1)[0] for k in weights_sd.keys()]))
+        remaining_keys.sort()
+        logger.info(f"Remaining LoRA modules after filtering: {remaining_keys}")
+        if len(weights_sd) == 0:
+            logger.warning("No keys left after filtering.")
+
+    return weights_sd
+
+
+def load_safetensors_with_lora_and_fp8(
+    model_files: Union[str, List[str]],
+    lora_weights_list: Optional[List[Dict[str, torch.Tensor]]],
+    lora_multipliers: Optional[List[float]],
+    fp8_optimization: bool,
+    calc_device: torch.device,
+    move_to_device: bool = False,
+    dit_weight_dtype: Optional[torch.dtype] = None,
+    target_keys: Optional[List[str]] = None,
+    exclude_keys: Optional[List[str]] = None,
+    disable_numpy_memmap: bool = False,
+    weight_transform_hooks: Optional[WeightTransformHooks] = None,
+    allow_prequantized_fp8: bool = False,
+) -> dict[str, torch.Tensor]:
+    """
+    Merge LoRA weights into the state dict of a model with fp8 optimization if needed.
+
+    Args:
+        model_files (Union[str, List[str]]): Path to the model file or list of paths. If the path matches a pattern like `00001-of-00004`, it will load all files with the same prefix.
+        lora_weights_list (Optional[List[Dict[str, torch.Tensor]]]): List of dictionaries of LoRA weight tensors to load.
+        lora_multipliers (Optional[List[float]]): List of multipliers for LoRA weights.
+        fp8_optimization (bool): Whether to apply FP8 optimization.
+        calc_device (torch.device): Device to calculate on.
+        move_to_device (bool): Whether to move tensors to the calculation device after loading.
+        target_keys (Optional[List[str]]): Keys to target for optimization.
+        exclude_keys (Optional[List[str]]): Keys to exclude from optimization.
+        disable_numpy_memmap (bool): Whether to disable numpy memmap when loading safetensors.
+        weight_transform_hooks (Optional[WeightTransformHooks]): Hooks for transforming weights during loading.
+    """
+
+    # if the file name ends with 00001-of-00004 etc, we need to load the files with the same prefix
+    if isinstance(model_files, str):
+        model_files = [model_files]
+
+    extended_model_files = []
+    for model_file in model_files:
+        split_filenames = get_split_weight_filenames(model_file)
+        if split_filenames is not None:
+            extended_model_files.extend(split_filenames)
+        else:
+            extended_model_files.append(model_file)
+    model_files = extended_model_files
+    logger.info(f"Loading model files: {model_files}")
+
+    # load LoRA weights
+    weight_hook = None
+    if lora_weights_list is None or len(lora_weights_list) == 0:
+        lora_weights_list = []
+        lora_multipliers = []
+        list_of_lora_weight_keys = []
+    else:
+        list_of_lora_weight_keys = []
+        for lora_sd in lora_weights_list:
+            lora_weight_keys = set(lora_sd.keys())
+            list_of_lora_weight_keys.append(lora_weight_keys)
+
+        if lora_multipliers is None:
+            lora_multipliers = [1.0] * len(lora_weights_list)
+        while len(lora_multipliers) < len(lora_weights_list):
+            lora_multipliers.append(1.0)
+        if len(lora_multipliers) > len(lora_weights_list):
+            lora_multipliers = lora_multipliers[: len(lora_weights_list)]
+
+        # detect network types for each lora_sd
+        lora_network_types = [detect_network_type(lora_sd) for lora_sd in lora_weights_list]
+
+        # Merge LoRA weights into the state dict
+        logger.info(f"Merging LoRA weights into state dict. multipliers: {lora_multipliers}, network types: {lora_network_types}")
+
+        # make hook for LoRA merging
+        def weight_hook_func(model_weight_key, model_weight: torch.Tensor, keep_on_calc_device=False):
+            nonlocal list_of_lora_weight_keys, lora_weights_list, lora_multipliers, lora_network_types, calc_device
+
+            if not model_weight_key.endswith(".weight"):
+                return model_weight
+
+            original_device = model_weight.device
+            if original_device != calc_device:
+                model_weight = model_weight.to(calc_device)  # to make calculation faster
+
+            for lora_weight_keys, lora_sd, multiplier, net_type in zip(
+                list_of_lora_weight_keys, lora_weights_list, lora_multipliers, lora_network_types
+            ):
+                lora_name = model_weight_key.rsplit(".", 1)[0]  # remove trailing ".weight"
+                lora_name = "lora_unet_" + lora_name.replace(".", "_")
+
+                if net_type == "loha":
+                    from musubi_tuner.networks.loha import merge_weights_to_tensor as loha_merge
+
+                    model_weight = loha_merge(model_weight, lora_name, lora_sd, lora_weight_keys, multiplier, calc_device)
+                elif net_type == "lokr":
+                    from musubi_tuner.networks.lokr import merge_weights_to_tensor as lokr_merge
+
+                    model_weight = lokr_merge(model_weight, lora_name, lora_sd, lora_weight_keys, multiplier, calc_device)
+                else:
+                    # standard LoRA (lora_down/lora_up)
+                    down_key = lora_name + ".lora_down.weight"
+                    up_key = lora_name + ".lora_up.weight"
+                    alpha_key = lora_name + ".alpha"
+                    if down_key not in lora_weight_keys or up_key not in lora_weight_keys:
+                        continue
+
+                    # get LoRA weights
+                    down_weight = lora_sd[down_key]
+                    up_weight = lora_sd[up_key]
+
+                    dim = down_weight.size()[0]
+                    alpha = lora_sd.get(alpha_key, dim)
+                    scale = alpha / dim
+
+                    down_weight = down_weight.to(calc_device)
+                    up_weight = up_weight.to(calc_device)
+
+                    original_dtype = model_weight.dtype
+                    if original_dtype.itemsize == 1:  # fp8
+                        # temporarily convert to float16 for calculation
+                        model_weight = model_weight.to(torch.float16)
+                        down_weight = down_weight.to(torch.float16)
+                        up_weight = up_weight.to(torch.float16)
+
+                    # W <- W + U * D
+                    if len(model_weight.size()) == 2:
+                        # linear
+                        if len(up_weight.size()) == 4:  # use linear projection mismatch
+                            up_weight = up_weight.squeeze(3).squeeze(2)
+                            down_weight = down_weight.squeeze(3).squeeze(2)
+                        model_weight = model_weight + multiplier * (up_weight @ down_weight) * scale
+                    elif down_weight.size()[2:4] == (1, 1):
+                        # conv2d 1x1
+                        model_weight = (
+                            model_weight
+                            + multiplier
+                            * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
+                            * scale
+                        )
+                    else:
+                        # conv2d 3x3
+                        conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+                        model_weight = model_weight + multiplier * conved * scale
+
+                    if original_dtype.itemsize == 1:  # fp8
+                        model_weight = model_weight.to(original_dtype)  # convert back to original dtype
+
+                    # remove LoRA keys from set
+                    lora_weight_keys.remove(down_key)
+                    lora_weight_keys.remove(up_key)
+                    if alpha_key in lora_weight_keys:
+                        lora_weight_keys.remove(alpha_key)
+
+            if not keep_on_calc_device and original_device != calc_device:
+                model_weight = model_weight.to(original_device)  # move back to original device
+            return model_weight
+
+        weight_hook = weight_hook_func
+
+    state_dict = load_safetensors_with_fp8_optimization_and_hook(
+        model_files,
+        fp8_optimization,
+        calc_device,
+        move_to_device,
+        dit_weight_dtype,
+        target_keys,
+        exclude_keys,
+        weight_hook=weight_hook,
+        disable_numpy_memmap=disable_numpy_memmap,
+        weight_transform_hooks=weight_transform_hooks,
+        allow_prequantized_fp8=allow_prequantized_fp8,
+    )
+
+    for lora_weight_keys in list_of_lora_weight_keys:
+        # check if all LoRA keys are used
+        if len(lora_weight_keys) > 0:
+            # if there are still LoRA keys left, it means they are not used in the model
+            # this is a warning, not an error
+            logger.warning(f"Warning: not all LoRA keys are used: {', '.join(lora_weight_keys)}")
+
+    return state_dict
+
+
+def load_safetensors_with_fp8_optimization_and_hook(
+    model_files: list[str],
+    fp8_optimization: bool,
+    calc_device: torch.device,
+    move_to_device: bool = False,
+    dit_weight_dtype: Optional[torch.dtype] = None,
+    target_keys: Optional[List[str]] = None,
+    exclude_keys: Optional[List[str]] = None,
+    weight_hook: callable = None,
+    disable_numpy_memmap: bool = False,
+    weight_transform_hooks: Optional[WeightTransformHooks] = None,
+    allow_prequantized_fp8: bool = False,
+) -> dict[str, torch.Tensor]:
+    """
+    Load state dict from safetensors files and merge LoRA weights into the state dict with fp8 optimization if needed.
+    """
+    if fp8_optimization:
+        logger.info(
+            f"Loading state dict with FP8 optimization. Dtype of weight: {dit_weight_dtype}, hook enabled: {weight_hook is not None}"
+        )
+        # dit_weight_dtype is not used because we use fp8 optimization
+        state_dict = load_safetensors_with_fp8_optimization(
+            model_files,
+            calc_device,
+            target_keys,
+            exclude_keys,
+            move_to_device=move_to_device,
+            weight_hook=weight_hook,
+            disable_numpy_memmap=disable_numpy_memmap,
+            weight_transform_hooks=weight_transform_hooks,
+            allow_prequantized_fp8=allow_prequantized_fp8,
+        )
+    else:
+        logger.info(
+            f"Loading state dict without FP8 optimization. Dtype of weight: {dit_weight_dtype}, hook enabled: {weight_hook is not None}"
+        )
+        state_dict = {}
+        for model_file in model_files:
+            with MemoryEfficientSafeOpen(model_file, disable_numpy_memmap=disable_numpy_memmap) as original_f:
+                f = TensorWeightAdapter(weight_transform_hooks, original_f) if weight_transform_hooks is not None else original_f
+                for key in tqdm(f.keys(), desc=f"Loading {os.path.basename(model_file)}", leave=False):
+                    if weight_hook is None and move_to_device:
+                        value = f.get_tensor(key, device=calc_device, dtype=dit_weight_dtype)
+                    else:
+                        value = f.get_tensor(key)  # we cannot directly load to device because get_tensor does non-blocking transfer
+                        if weight_hook is not None:
+                            value = weight_hook(key, value, keep_on_calc_device=move_to_device)
+                        if move_to_device:
+                            value = value.to(calc_device, dtype=dit_weight_dtype, non_blocking=True)
+                        elif dit_weight_dtype is not None:
+                            value = value.to(dit_weight_dtype)
+
+                    state_dict[key] = value
+        if move_to_device:
+            synchronize_device(calc_device)
+
+    return state_dict
diff --git a/musubi_tuner/utils/model_utils.py b/musubi_tuner/utils/model_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..78f6fa3ed02ca8ea71734e1eae2f9b17dcd43638
--- /dev/null
+++ b/musubi_tuner/utils/model_utils.py
@@ -0,0 +1,295 @@
+import argparse
+import hashlib
+from functools import lru_cache
+from io import BytesIO
+from typing import Any, Callable, Optional
+import logging
+import safetensors.torch
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def model_hash(filename):
+    """Old model hash used by stable-diffusion-webui"""
+    try:
+        with open(filename, "rb") as file:
+            m = hashlib.sha256()
+
+            file.seek(0x100000)
+            m.update(file.read(0x10000))
+            return m.hexdigest()[0:8]
+    except FileNotFoundError:
+        return "NOFILE"
+    except IsADirectoryError:  # Linux?
+        return "IsADirectory"
+    except PermissionError:  # Windows
+        return "IsADirectory"
+
+
+def calculate_sha256(filename):
+    """New model hash used by stable-diffusion-webui"""
+    try:
+        hash_sha256 = hashlib.sha256()
+        blksize = 1024 * 1024
+
+        with open(filename, "rb") as f:
+            for chunk in iter(lambda: f.read(blksize), b""):
+                hash_sha256.update(chunk)
+
+        return hash_sha256.hexdigest()
+    except FileNotFoundError:
+        return "NOFILE"
+    except IsADirectoryError:  # Linux?
+        return "IsADirectory"
+    except PermissionError:  # Windows
+        return "IsADirectory"
+
+
+def addnet_hash_legacy(b):
+    """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+    m = hashlib.sha256()
+
+    b.seek(0x100000)
+    m.update(b.read(0x10000))
+    return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+    """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+    hash_sha256 = hashlib.sha256()
+    blksize = 1024 * 1024
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+    for chunk in iter(lambda: b.read(blksize), b""):
+        hash_sha256.update(chunk)
+
+    return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+    """Precalculate the model hashes needed by sd-webui-additional-networks to
+    save time on indexing the model later."""
+
+    # Because writing user metadata to the file can change the result of
+    # sd_models.model_hash(), only retain the training metadata for purposes of
+    # calculating the hash, as they are meant to be immutable
+    metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+    bytes = safetensors.torch.save(tensors, metadata)
+    b = BytesIO(bytes)
+
+    model_hash = addnet_hash_safetensors(b)
+    legacy_hash = addnet_hash_legacy(b)
+    return model_hash, legacy_hash
+
+
+def dtype_to_str(dtype: torch.dtype) -> str:
+    # get name of the dtype
+    dtype_name = str(dtype).split(".")[-1]
+    return dtype_name
+
+
+def str_to_dtype(s: Optional[str], default_dtype: Optional[torch.dtype] = None) -> torch.dtype:
+    """
+    Convert a string to a torch.dtype
+
+    Args:
+        s: string representation of the dtype
+        default_dtype: default dtype to return if s is None
+
+    Returns:
+        torch.dtype: the corresponding torch.dtype
+
+    Raises:
+        ValueError: if the dtype is not supported
+
+    Examples:
+        >>> str_to_dtype("float32")
+        torch.float32
+        >>> str_to_dtype("fp32")
+        torch.float32
+        >>> str_to_dtype("float16")
+        torch.float16
+        >>> str_to_dtype("fp16")
+        torch.float16
+        >>> str_to_dtype("bfloat16")
+        torch.bfloat16
+        >>> str_to_dtype("bf16")
+        torch.bfloat16
+        >>> str_to_dtype("fp8")
+        torch.float8_e4m3fn
+        >>> str_to_dtype("fp8_e4m3fn")
+        torch.float8_e4m3fn
+        >>> str_to_dtype("fp8_e4m3fnuz")
+        torch.float8_e4m3fnuz
+        >>> str_to_dtype("fp8_e5m2")
+        torch.float8_e5m2
+        >>> str_to_dtype("fp8_e5m2fnuz")
+        torch.float8_e5m2fnuz
+    """
+    if s is None:
+        return default_dtype
+    if s in ["bf16", "bfloat16"]:
+        return torch.bfloat16
+    elif s in ["fp16", "float16"]:
+        return torch.float16
+    elif s in ["fp32", "float32", "float"]:
+        return torch.float32
+    elif s in ["fp8_e4m3fn", "e4m3fn", "float8_e4m3fn"]:
+        return torch.float8_e4m3fn
+    elif s in ["fp8_e4m3fnuz", "e4m3fnuz", "float8_e4m3fnuz"]:
+        return torch.float8_e4m3fnuz
+    elif s in ["fp8_e5m2", "e5m2", "float8_e5m2"]:
+        return torch.float8_e5m2
+    elif s in ["fp8_e5m2fnuz", "e5m2fnuz", "float8_e5m2fnuz"]:
+        return torch.float8_e5m2fnuz
+    elif s in ["fp8", "float8"]:
+        return torch.float8_e4m3fn  # default fp8
+    else:
+        raise ValueError(f"Unsupported dtype: {s}")
+
+
+@lru_cache(maxsize=1)
+def _known_dtype_strs() -> tuple[str, ...]:
+    """All dtype strings dtype_to_str() can emit, longest first.
+
+    Sorting by length keeps multi-underscore names (e.g. "float8_e4m3fn") ahead of their
+    prefixes ("float8...") so the longest matching suffix wins.
+    """
+    names = set()
+    for attr in dir(torch):
+        try:
+            obj = getattr(torch, attr)
+        except Exception:
+            continue
+        if isinstance(obj, torch.dtype):
+            names.add(dtype_to_str(obj))
+    return tuple(sorted(names, key=len, reverse=True))
+
+
+def remove_dtype_suffix(name: str) -> str:
+    """Remove a trailing ``_<dtype>`` suffix (as written by dtype_to_str) from a cache key.
+
+    Robust to dtype names that contain underscores such as ``float8_e4m3fn``; a plain
+    ``rsplit("_", 1)`` would only drop the final ``fn`` segment. Returns ``name`` unchanged
+    if it does not end with a known dtype suffix.
+    """
+    for dtype_str in _known_dtype_strs():
+        suffix = "_" + dtype_str
+        if name.endswith(suffix):
+            return name[: -len(suffix)]
+    return name
+
+
+def to_device(x: Any, device: torch.device) -> Any:
+    if isinstance(x, torch.Tensor):
+        return x.to(device)
+    elif isinstance(x, list):
+        return [to_device(elem, device) for elem in x]
+    elif isinstance(x, tuple):
+        return tuple(to_device(elem, device) for elem in x)
+    elif isinstance(x, dict):
+        return {k: to_device(v, device) for k, v in x.items()}
+    else:
+        return x
+
+
+def to_cpu(x: Any) -> Any:
+    """
+    Recursively moves torch.Tensor objects (and containers thereof) to CPU.
+
+    Args:
+        x: A torch.Tensor, or a (possibly nested) list, tuple, or dict containing tensors.
+
+    Returns:
+        The same structure as x, with all torch.Tensor objects moved to CPU.
+        Non-tensor objects are returned unchanged.
+    """
+    if isinstance(x, torch.Tensor):
+        return x.cpu()
+    elif isinstance(x, list):
+        return [to_cpu(elem) for elem in x]
+    elif isinstance(x, tuple):
+        return tuple(to_cpu(elem) for elem in x)
+    elif isinstance(x, dict):
+        return {k: to_cpu(v) for k, v in x.items()}
+    else:
+        return x
+
+
+def create_cpu_offloading_wrapper(func: Callable, device: torch.device) -> Callable:
+    """
+    Create a wrapper function that offloads inputs to CPU before calling the original function
+    and moves outputs back to the specified device.
+
+    Args:
+        func: The original function to wrap.
+        device: The device to move outputs back to.
+
+    Returns:
+        A wrapped function that offloads inputs to CPU and moves outputs back to the specified device.
+    """
+
+    def wrapper(orig_func: Callable) -> Callable:
+        def custom_forward(*inputs):
+            nonlocal device, orig_func
+            cuda_inputs = to_device(inputs, device)
+            outputs = orig_func(*cuda_inputs)
+            return to_cpu(outputs)
+
+        return custom_forward
+
+    return wrapper(func)
+
+
+def disable_linear_from_compile(module: torch.nn.Module):
+    """Monkey-patch to disable torch.compile for all Linear layers (if the class name ends with 'Linear') in the given module."""
+    for sub_module in module.modules():
+        # if isinstance(sub_module, torch.nn.Linear):
+        if sub_module.__class__.__name__.endswith("Linear"):
+            if not hasattr(sub_module, "_forward_before_disable_compile"):
+                sub_module._forward_before_disable_compile = sub_module.forward
+                sub_module._eager_forward = torch._dynamo.disable()(sub_module.forward)
+            sub_module.forward = sub_module._eager_forward  # override forward to disable compile
+
+
+def compile_transformer(
+    args: argparse.Namespace,
+    transformer: torch.nn.Module,
+    target_blocks: list[torch.nn.ModuleList | list[torch.nn.Module]],
+    disable_linear: bool,
+) -> torch.nn.Module:
+    if disable_linear:
+        logger.info("Disable linear from torch.compile for swap blocks...")
+        for blocks in target_blocks:
+            for block in blocks:
+                disable_linear_from_compile(block)
+
+    compile_dynamic = None
+    if args.compile_dynamic is not None:
+        compile_dynamic = {"true": True, "false": False, "auto": None}[args.compile_dynamic.lower()]
+
+    logger.info(
+        f"Compiling DiT model with torch.compile: backend={args.compile_backend}, mode={args.compile_mode}, dynamic={compile_dynamic}, fullgraph={args.compile_fullgraph}"
+    )
+
+    if args.compile_cache_size_limit is not None:
+        torch._dynamo.config.cache_size_limit = args.compile_cache_size_limit
+
+    for blocks in target_blocks:
+        for i, block in enumerate(blocks):
+            block = torch.compile(
+                block,
+                backend=args.compile_backend,
+                mode=args.compile_mode,
+                dynamic=compile_dynamic,
+                fullgraph=args.compile_fullgraph,
+            )
+            blocks[i] = block
+    return transformer
diff --git a/musubi_tuner/utils/safetensors_utils.py b/musubi_tuner/utils/safetensors_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b50939954d701acc655d9b10fa89b61ce75471bf
--- /dev/null
+++ b/musubi_tuner/utils/safetensors_utils.py
@@ -0,0 +1,483 @@
+from dataclasses import dataclass
+import os
+import re
+import numpy as np
+import torch
+import json
+import struct
+from typing import Dict, Any, Union, Optional
+
+from safetensors.torch import load_file
+
+from musubi_tuner.utils.device_utils import synchronize_device
+
+
+def mem_eff_save_file(tensors: Dict[str, torch.Tensor], filename: str, metadata: Dict[str, Any] = None):
+    """
+    memory efficient save file
+    """
+
+    _TYPES = {
+        torch.float64: "F64",
+        torch.float32: "F32",
+        torch.float16: "F16",
+        torch.bfloat16: "BF16",
+        torch.int64: "I64",
+        torch.int32: "I32",
+        torch.int16: "I16",
+        torch.int8: "I8",
+        torch.uint8: "U8",
+        torch.bool: "BOOL",
+        getattr(torch, "float8_e5m2", None): "F8_E5M2",
+        getattr(torch, "float8_e4m3fn", None): "F8_E4M3",
+    }
+    _ALIGN = 256
+
+    def validate_metadata(metadata: Dict[str, Any]) -> Dict[str, str]:
+        validated = {}
+        for key, value in metadata.items():
+            if not isinstance(key, str):
+                raise ValueError(f"Metadata key must be a string, got {type(key)}")
+            if not isinstance(value, str):
+                print(f"Warning: Metadata value for key '{key}' is not a string. Converting to string.")
+                validated[key] = str(value)
+            else:
+                validated[key] = value
+        return validated
+
+    # print(f"Using memory efficient save file: {filename}")
+
+    header = {}
+    offset = 0
+    if metadata:
+        header["__metadata__"] = validate_metadata(metadata)
+    for k, v in tensors.items():
+        if v.numel() == 0:  # empty tensor
+            header[k] = {"dtype": _TYPES[v.dtype], "shape": list(v.shape), "data_offsets": [offset, offset]}
+        else:
+            size = v.numel() * v.element_size()
+            header[k] = {"dtype": _TYPES[v.dtype], "shape": list(v.shape), "data_offsets": [offset, offset + size]}
+            offset += size
+
+    hjson = json.dumps(header).encode("utf-8")
+    hjson += b" " * (-(len(hjson) + 8) % _ALIGN)
+
+    with open(filename, "wb") as f:
+        f.write(struct.pack("<Q", len(hjson)))
+        f.write(hjson)
+
+        for k, v in tensors.items():
+            if v.numel() == 0:
+                continue
+            if v.is_cuda:
+                # Direct GPU to disk save
+                with torch.cuda.device(v.device):
+                    if v.dim() == 0:  # if scalar, need to add a dimension to work with view
+                        v = v.unsqueeze(0)
+                    tensor_bytes = v.contiguous().view(torch.uint8)
+                    tensor_bytes.cpu().numpy().tofile(f)
+            else:
+                # CPU tensor save
+                if v.dim() == 0:  # if scalar, need to add a dimension to work with view
+                    v = v.unsqueeze(0)
+                v.contiguous().view(torch.uint8).numpy().tofile(f)
+
+
+class MemoryEfficientSafeOpen:
+    """Memory-efficient reader for safetensors files.
+
+    This class provides a memory-efficient way to read tensors from safetensors files
+    by using memory mapping for large tensors and avoiding unnecessary copies.
+    """
+
+    def __init__(self, filename, disable_numpy_memmap=False):
+        """Initialize the SafeTensor reader.
+
+        Args:
+            filename (str): Path to the safetensors file to read.
+            disable_numpy_memmap (bool): If True, disable numpy memory mapping for large tensors, using standard file read instead.
+        """
+        self.filename = filename
+        self.file = open(filename, "rb")
+        self.header, self.header_size = self._read_header()
+        self.disable_numpy_memmap = disable_numpy_memmap
+
+    def __enter__(self):
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and close file."""
+        self.file.close()
+
+    def keys(self):
+        """Get all tensor keys in the file.
+
+        Returns:
+            list: List of tensor names (excludes metadata).
+        """
+        return [k for k in self.header.keys() if k != "__metadata__"]
+
+    def metadata(self) -> Dict[str, str]:
+        """Get metadata from the file.
+
+        Returns:
+            Dict[str, str]: Metadata dictionary.
+        """
+        return self.header.get("__metadata__", {})
+
+    def _read_header(self):
+        """Read and parse the header from the safetensors file.
+
+        Returns:
+            tuple: (header_dict, header_size) containing parsed header and its size.
+        """
+        # Read header size (8 bytes, little-endian unsigned long long)
+        header_size = struct.unpack("<Q", self.file.read(8))[0]
+        # Read and decode header JSON
+        header_json = self.file.read(header_size).decode("utf-8")
+        return json.loads(header_json), header_size
+
+    def get_tensor(self, key: str, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        """Load a tensor from the file with memory-efficient strategies.
+
+        **Note:**
+        If device is 'cuda' , the transfer to GPU is done efficiently using pinned memory and non-blocking transfer.
+        So you must ensure that the transfer is completed before using the tensor (e.g., by `torch.cuda.synchronize()`).
+
+        If the tensor is large (>10MB) and the target device is CUDA, memory mapping with numpy.memmap is used to avoid intermediate copies.
+
+        Args:
+            key (str): Name of the tensor to load.
+            device (Optional[torch.device]): Target device for the tensor.
+            dtype (Optional[torch.dtype]): Target dtype for the tensor.
+
+        Returns:
+            torch.Tensor: The loaded tensor.
+
+        Raises:
+            KeyError: If the tensor key is not found in the file.
+        """
+        if key not in self.header:
+            raise KeyError(f"Tensor '{key}' not found in the file")
+
+        metadata = self.header[key]
+        offset_start, offset_end = metadata["data_offsets"]
+        num_bytes = offset_end - offset_start
+
+        original_dtype = self._get_torch_dtype(metadata["dtype"])
+        target_dtype = dtype if dtype is not None else original_dtype
+
+        # Handle empty tensors
+        if num_bytes == 0:
+            return torch.empty(metadata["shape"], dtype=target_dtype, device=device)
+
+        # Determine if we should use pinned memory for GPU transfer
+        non_blocking = device is not None and device.type == "cuda"
+
+        # Calculate absolute file offset
+        tensor_offset = self.header_size + 8 + offset_start  # adjust offset by header size
+
+        # Memory mapping strategy for large tensors to GPU
+        # Use memmap for large tensors to avoid intermediate copies.
+        # If device is cpu, tensor is not copied to gpu, so using memmap locks the file, which is not desired.
+        # So we only use memmap if device is not cpu.
+        # If disable_numpy_memmap is True, skip numpy memory mapping to load with standard file read.
+        if not self.disable_numpy_memmap and num_bytes > 10 * 1024 * 1024 and device is not None and device.type != "cpu":
+            # Create memory map for zero-copy reading
+            mm = np.memmap(self.filename, mode="c", dtype=np.uint8, offset=tensor_offset, shape=(num_bytes,))
+            byte_tensor = torch.from_numpy(mm)  # zero copy
+            del mm
+
+            # Deserialize tensor (view and reshape)
+            cpu_tensor = self._deserialize_tensor(byte_tensor, metadata)  # view and reshape
+            del byte_tensor
+
+            # Transfer to target device and dtype
+            gpu_tensor = cpu_tensor.to(device=device, dtype=target_dtype, non_blocking=non_blocking)
+            del cpu_tensor
+            return gpu_tensor
+
+        # Standard file reading strategy for smaller tensors or CPU target
+        # seek to the specified position
+        self.file.seek(tensor_offset)
+
+        # read directly into a numpy array by numpy.fromfile without intermediate copy
+        numpy_array = np.fromfile(self.file, dtype=np.uint8, count=num_bytes)
+        byte_tensor = torch.from_numpy(numpy_array)
+        del numpy_array
+
+        # deserialize (view and reshape)
+        deserialized_tensor = self._deserialize_tensor(byte_tensor, metadata)
+        del byte_tensor
+
+        # cast to target dtype and move to device
+        return deserialized_tensor.to(device=device, dtype=target_dtype, non_blocking=non_blocking)
+
+    def _deserialize_tensor(self, byte_tensor: torch.Tensor, metadata: Dict):
+        """Deserialize byte tensor to the correct shape and dtype.
+
+        Args:
+            byte_tensor (torch.Tensor): Raw byte tensor from file.
+            metadata (Dict): Tensor metadata containing dtype and shape info.
+
+        Returns:
+            torch.Tensor: Deserialized tensor with correct shape and dtype.
+        """
+        dtype = self._get_torch_dtype(metadata["dtype"])
+        shape = metadata["shape"]
+
+        # Handle special float8 types
+        if metadata["dtype"] in ["F8_E5M2", "F8_E4M3"]:
+            return self._convert_float8(byte_tensor, metadata["dtype"], shape)
+
+        # Standard conversion: view as target dtype and reshape
+        return byte_tensor.view(dtype).reshape(shape)
+
+    @staticmethod
+    def _get_torch_dtype(dtype_str):
+        """Convert string dtype to PyTorch dtype.
+
+        Args:
+            dtype_str (str): String representation of the dtype.
+
+        Returns:
+            torch.dtype: Corresponding PyTorch dtype.
+        """
+        # Standard dtype mappings
+        dtype_map = {
+            "F64": torch.float64,
+            "F32": torch.float32,
+            "F16": torch.float16,
+            "BF16": torch.bfloat16,
+            "I64": torch.int64,
+            "I32": torch.int32,
+            "I16": torch.int16,
+            "I8": torch.int8,
+            "U8": torch.uint8,
+            "BOOL": torch.bool,
+        }
+        # Add float8 types if available in PyTorch version
+        if hasattr(torch, "float8_e5m2"):
+            dtype_map["F8_E5M2"] = torch.float8_e5m2
+        if hasattr(torch, "float8_e4m3fn"):
+            dtype_map["F8_E4M3"] = torch.float8_e4m3fn
+        return dtype_map.get(dtype_str)
+
+    @staticmethod
+    def _convert_float8(byte_tensor, dtype_str, shape):
+        """Convert byte tensor to float8 format if supported.
+
+        Args:
+            byte_tensor (torch.Tensor): Raw byte tensor.
+            dtype_str (str): Float8 dtype string ("F8_E5M2" or "F8_E4M3").
+            shape (tuple): Target tensor shape.
+
+        Returns:
+            torch.Tensor: Tensor with float8 dtype.
+
+        Raises:
+            ValueError: If float8 type is not supported in current PyTorch version.
+        """
+        # Convert to specific float8 types if available
+        if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
+            return byte_tensor.view(torch.float8_e5m2).reshape(shape)
+        elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
+            return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
+        else:
+            # Float8 not supported in this PyTorch version
+            raise ValueError(f"Unsupported float8 type: {dtype_str} (upgrade PyTorch to support float8 types)")
+
+
+def load_safetensors(
+    path: str,
+    device: Union[str, torch.device],
+    disable_mmap: bool = False,
+    dtype: Optional[torch.dtype] = None,
+    disable_numpy_memmap: bool = False,
+) -> dict[str, torch.Tensor]:
+    if disable_mmap:
+        # return safetensors.torch.load(open(path, "rb").read())
+        # use experimental loader
+        # logger.info(f"Loading without mmap (experimental)")
+        state_dict = {}
+        device = torch.device(device) if device is not None else None
+        with MemoryEfficientSafeOpen(path, disable_numpy_memmap=disable_numpy_memmap) as f:
+            for key in f.keys():
+                state_dict[key] = f.get_tensor(key, device=device, dtype=dtype)
+        synchronize_device(device)
+        return state_dict
+    else:
+        try:
+            state_dict = load_file(path, device=device)
+        except:
+            state_dict = load_file(path)  # prevent device invalid Error
+        if dtype is not None:
+            for key in state_dict.keys():
+                state_dict[key] = state_dict[key].to(dtype=dtype)
+        return state_dict
+
+
+def get_split_weight_filenames(file_path: str) -> Optional[list[str]]:
+    """
+    Get the list of split weight filenames (full paths) if the file name ends with 00001-of-00004 etc.
+    Returns None if the file is not split.
+    """
+    basename = os.path.basename(file_path)
+    match = re.match(r"^(.*?)(\d+)-of-(\d+)\.safetensors$", basename)
+    if match:
+        prefix = basename[: match.start(2)]
+        count = int(match.group(3))
+        filenames = []
+        for i in range(count):
+            filename = f"{prefix}{i + 1:05d}-of-{count:05d}.safetensors"
+            filepath = os.path.join(os.path.dirname(file_path), filename)
+            if os.path.exists(filepath):
+                filenames.append(filepath)
+            else:
+                raise FileNotFoundError(f"File {filepath} not found")
+        return filenames
+    else:
+        return None
+
+
+def load_split_weights(
+    file_path: str, device: Union[str, torch.device] = "cpu", disable_mmap: bool = False, dtype: Optional[torch.dtype] = None
+) -> Dict[str, torch.Tensor]:
+    """
+    Load split weights from a file. If the file name ends with 00001-of-00004 etc, it will load all files with the same prefix.
+    dtype is as is, no conversion is done.
+    """
+    device = torch.device(device)
+
+    # if the file name ends with 00001-of-00004 etc, we need to load the files with the same prefix
+    split_filenames = get_split_weight_filenames(file_path)
+    if split_filenames is not None:
+        state_dict = {}
+        for filename in split_filenames:
+            state_dict.update(load_safetensors(filename, device=device, disable_mmap=disable_mmap, dtype=dtype))
+    else:
+        state_dict = load_safetensors(file_path, device=device, disable_mmap=disable_mmap, dtype=dtype)
+    return state_dict
+
+
+def find_key(safetensors_file: str, starts_with: Optional[str] = None, ends_with: Optional[str] = None) -> Optional[str]:
+    """
+    Find a key in a safetensors file that starts with `starts_with` and ends with `ends_with`.
+    If `starts_with` is None, it will match any key.
+    If `ends_with` is None, it will match any key.
+    Returns the first matching key or None if no key matches.
+    """
+    with MemoryEfficientSafeOpen(safetensors_file) as f:
+        for key in f.keys():
+            if (starts_with is None or key.startswith(starts_with)) and (ends_with is None or key.endswith(ends_with)):
+                return key
+    return None
+
+
+def find_keys(safetensors_file: str, starts_with: Optional[str] = None, ends_with: Optional[str] = None) -> list[str]:
+    """
+    Find all keys in a safetensors file that start with `starts_with` and end with `ends_with`.
+    If `starts_with` is None, it will match any key. If `ends_with` is None, it will match any key.
+
+    The matching keys are returned sorted, so callers get a deterministic order regardless of the
+    order the keys happen to be stored in the file header. This matters when the keys are folded into
+    a bucket key (e.g. ``latents_control_{i}``): otherwise two files holding the same tensors in a
+    different header order, or with the index<->shape association swapped, could end up in different
+    buckets or be wrongly batched together.
+    """
+    with MemoryEfficientSafeOpen(safetensors_file) as f:
+        keys = [
+            key
+            for key in f.keys()
+            if (starts_with is None or key.startswith(starts_with)) and (ends_with is None or key.endswith(ends_with))
+        ]
+    return sorted(keys)
+
+
+@dataclass
+class WeightTransformHooks:
+    split_hook: Optional[callable] = None
+    concat_hook: Optional[callable] = None
+
+
+class TensorWeightAdapter:
+    """
+    A wrapper for weight conversion hooks (split and concat) to be used with MemoryEfficientSafeOpen.
+    This wrapper adapts the original MemoryEfficientSafeOpen to apply the provided split and concat hooks
+    when loading tensors.
+
+    split_hook: A callable that takes (original_key: str, original_tensor: torch.Tensor) and returns (new_keys: list[str], new_tensors: list[torch.Tensor]).
+    concat_hook: A callable that takes (original_key: str, tensors: dict[str, torch.Tensor]) and returns (new_key: str,  concatenated_tensor: torch.Tensor).
+
+    If tensors is None, the hook should return only the new keys (for split) or new key (for concat), without tensors.
+
+    No need to implement __enter__ and __exit__ methods, as they are handled by the original MemoryEfficientSafeOpen.
+    Do not use this wrapper as a context manager directly, like `with WeightConvertHookWrapper(...) as f:`.
+
+    **concat_hook is not tested yet.**
+    """
+
+    def __init__(self, weight_convert_hook: WeightTransformHooks, original_f: MemoryEfficientSafeOpen):
+        self.original_f = original_f
+        self.new_key_to_original_key_map: dict[
+            str, Union[str, list[str]]
+        ] = {}  # for split: new_key -> original_key; for concat: new_key -> list of original_keys
+        self.concat_key_set = set()  # set of keys that are created by concat_hook, to distinguish from split_hook
+        self.new_keys = []
+        self.tensor_cache = {}  # cache for split tensors
+        self.split_hook = weight_convert_hook.split_hook
+        self.concat_hook = weight_convert_hook.concat_hook
+
+        for key in self.original_f.keys():
+            if self.split_hook is not None:
+                converted_keys, _ = self.split_hook(key, None)  # get new keys only
+                if converted_keys is not None:
+                    for new_key in converted_keys:
+                        self.new_key_to_original_key_map[new_key] = key
+                    self.new_keys.extend(converted_keys)
+                    continue  # skip concat_hook if split_hook is applied
+
+            if self.concat_hook is not None:
+                converted_key, _ = self.concat_hook(key, None)  # get new key only
+                if converted_key is not None:
+                    if converted_key not in self.concat_key_set:  # first time seeing this concatenated key
+                        self.concat_key_set.add(converted_key)
+                        self.new_key_to_original_key_map[converted_key] = []
+
+                    # multiple original keys map to the same concatenated key
+                    self.new_key_to_original_key_map[converted_key].append(key)
+
+                    self.new_keys.append(converted_key)
+                    continue  # skip to next key
+
+            # direct mapping
+            self.new_keys.append(key)
+
+    def keys(self) -> list[str]:
+        return self.new_keys
+
+    def get_tensor(self, new_key: str, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        # load tensor by new_key, applying split or concat hooks as needed
+        if new_key not in self.new_key_to_original_key_map:
+            # direct mapping
+            return self.original_f.get_tensor(new_key, device=device, dtype=dtype)
+
+        elif new_key not in self.concat_key_set:
+            # split hook: split key is requested multiple times, so we cache the result
+            original_key = self.new_key_to_original_key_map[new_key]
+            if original_key not in self.tensor_cache:  # not yet split
+                original_tensor = self.original_f.get_tensor(original_key, device=device, dtype=dtype)
+                new_keys, new_tensors = self.split_hook(original_key, original_tensor)  # apply split hook
+                for k, t in zip(new_keys, new_tensors):
+                    self.tensor_cache[k] = t
+            return self.tensor_cache.pop(new_key)  # return and remove from cache
+
+        else:
+            # concat hook: concatenated key is requested only once, so we do not cache the result
+            tensors = {}
+            for original_key in self.new_key_to_original_key_map[new_key]:
+                tensor = self.original_f.get_tensor(original_key, device=device, dtype=dtype)
+                tensors[original_key] = tensor
+            _, concatenated_tensors = self.concat_hook(self.new_key_to_original_key_map[new_key][0], tensors)  # apply concat hook
+            return concatenated_tensors
diff --git a/musubi_tuner/utils/sai_model_spec.py b/musubi_tuner/utils/sai_model_spec.py
new file mode 100644
index 0000000000000000000000000000000000000000..01b155583aae7b387d54ac6021f46d4cfdb80bb3
--- /dev/null
+++ b/musubi_tuner/utils/sai_model_spec.py
@@ -0,0 +1,403 @@
+# based on https://github.com/Stability-AI/ModelSpec
+import datetime
+import hashlib
+from io import BytesIO
+import os
+from typing import List, Optional, Tuple, Union
+import safetensors
+import logging
+
+from musubi_tuner.dataset.image_video_dataset import (
+    ARCHITECTURE_HUNYUAN_VIDEO,
+    ARCHITECTURE_HUNYUAN_VIDEO_1_5,
+    ARCHITECTURE_HIDREAM_O1,
+    ARCHITECTURE_IDEOGRAM4,
+    ARCHITECTURE_QWEN_IMAGE,
+    ARCHITECTURE_QWEN_IMAGE_EDIT,
+    ARCHITECTURE_QWEN_IMAGE_LAYERED,
+    ARCHITECTURE_WAN,
+    ARCHITECTURE_FRAMEPACK,
+    ARCHITECTURE_FLUX_KONTEXT,
+    ARCHITECTURE_FLUX_2_DEV,
+    ARCHITECTURE_FLUX_2_KLEIN_4B,
+    ARCHITECTURE_FLUX_2_KLEIN_9B,
+    ARCHITECTURE_KANDINSKY5,
+    ARCHITECTURE_KREA2,
+    ARCHITECTURE_Z_IMAGE,
+)
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+r"""
+# Metadata Example
+metadata = {
+    # === Must ===
+    "modelspec.sai_model_spec": "1.0.0", # Required version ID for the spec
+    "modelspec.architecture": "stable-diffusion-xl-v1-base", # Architecture, reference the ID of the original model of the arch to match the ID
+    "modelspec.implementation": "sgm",
+    "modelspec.title": "Example Model Version 1.0", # Clean, human-readable title. May use your own phrasing/language/etc
+    # === Should ===
+    "modelspec.author": "Example Corp", # Your name or company name
+    "modelspec.description": "This is my example model to show you how to do it!", # Describe the model in your own words/language/etc. Focus on what users need to know
+    "modelspec.date": "2023-07-20", # ISO-8601 compliant date of when the model was created
+    # === Can ===
+    "modelspec.license": "ExampleLicense-1.0", # eg CreativeML Open RAIL, etc.
+    "modelspec.usage_hint": "Use keyword 'example'" # In your own language, very short hints about how the user should use the model
+}
+"""
+
+BASE_METADATA = {
+    # === Must ===
+    "modelspec.sai_model_spec": "1.0.0",  # Required version ID for the spec
+    "modelspec.architecture": None,
+    "modelspec.implementation": None,
+    "modelspec.title": None,
+    "modelspec.resolution": None,
+    # === Should ===
+    "modelspec.description": None,
+    "modelspec.author": None,
+    "modelspec.date": None,
+    # === Can ===
+    "modelspec.license": None,
+    "modelspec.tags": None,
+    "modelspec.merged_from": None,
+    "modelspec.prediction_type": None,
+    "modelspec.timestep_range": None,
+    "modelspec.encoder_layer": None,
+}
+
+# 別に使うやつだけ定義
+MODELSPEC_TITLE = "modelspec.title"
+
+ARCH_HUNYUAN_VIDEO = "hunyuan-video"
+
+# Official Wan2.1 weights does not have sai_model_spec, so we use this as an architecture name
+ARCH_WAN = "wan2.1"
+
+ARCH_FRAMEPACK = "framepack"
+ARCH_FLUX_KONTEXT = "Flux.1-dev"
+ARCH_FLUX_2_DEV = "Flux.2-dev"
+ARCH_FLUX_2_KLEIN_4B = "Flux.2-klein-4b"
+ARCH_FLUX_2_KLEIN_9B = "Flux.2-klein-9b"
+ARCH_QWEN_IMAGE = "Qwen-Image"
+ARCH_QWEN_IMAGE_EDIT = "Qwen-Image-Edit"
+ARCH_QWEN_IMAGE_EDIT_PLUS = "Qwen-Image-Edit-Plus"
+ARCH_QWEN_IMAGE_EDIT_2511 = "Qwen-Image-Edit-2511"
+CUSTOM_ARCH_QWEN_IMAGE_EDIT_PLUS = "@@Qwen-Image-Edit-Plus@@"  # special custom architecture name for Qwen-Image-Edit-Plus
+CUSTOM_ARCH_QWEN_IMAGE_EDIT_2511 = "@@Qwen-Image-Edit-2511@@"  # special custom architecture name for Qwen-Image-Edit-2511
+ARCH_QWEN_IMAGE_LAYERED = "Qwen-Image-Layered"
+ARCH_KANDINSKY5 = "Kandinsky-5"
+ARCH_HUNYUAN_VIDEO_1_5 = "hunyuan-video-1.5"
+ARCH_Z_IMAGE = "Z-Image"
+ARCH_HIDREAM_O1 = "HiDream-O1-Image"
+ARCH_IDEOGRAM4 = "Ideogram-4"
+ARCH_KREA2 = "Krea-2"
+
+ADAPTER_LORA = "lora"
+
+IMPL_HUNYUAN_VIDEO = "https://github.com/Tencent/HunyuanVideo"
+IMPL_WAN = "https://github.com/Wan-Video/Wan2.1"
+IMPL_FRAMEPACK = "https://github.com/lllyasviel/FramePack"
+IMPL_FLUX_KONTEXT = "https://github.com/black-forest-labs/flux"
+IMPL_FLUX_2 = "https://github.com/black-forest-labs/flux2"
+IMPL_QWEN_IMAGE = "https://github.com/QwenLM/Qwen-Image"
+IMPL_QWEN_IMAGE_EDIT = IMPL_QWEN_IMAGE
+IMPL_QWEN_IMAGE_LAYERED = "https://github.com/QwenLM/Qwen-Image-Layered"
+IMPL_KANDINSKY5 = "https://github.com/kandinskylab/kandinsky-5"
+IMPL_HUNYUAN_VIDEO_1_5 = "https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5"
+IMPL_Z_IMAGE = "https://github.com/Tongyi-MAI/Z-Image"
+IMPL_HIDREAM_O1 = "https://github.com/HiDream-ai/HiDream-O1-Image"
+IMPL_IDEOGRAM4 = "https://huggingface.co/Comfy-Org/Ideogram-4"
+IMPL_KREA2 = "https://github.com/krea-ai/krea-2"
+
+PRED_TYPE_EPSILON = "epsilon"
+# PRED_TYPE_V = "v"
+
+
+def load_bytes_in_safetensors(tensors):
+    bytes = safetensors.torch.save(tensors)
+    b = BytesIO(bytes)
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+
+    return b.read()
+
+
+def precalculate_safetensors_hashes(state_dict):
+    # calculate each tensor one by one to reduce memory usage
+    hash_sha256 = hashlib.sha256()
+    for tensor in state_dict.values():
+        single_tensor_sd = {"tensor": tensor}
+        bytes_for_tensor = load_bytes_in_safetensors(single_tensor_sd)
+        hash_sha256.update(bytes_for_tensor)
+
+    return f"0x{hash_sha256.hexdigest()}"
+
+
+def update_hash_sha256(metadata: dict, state_dict: dict):
+    raise NotImplementedError
+
+
+def build_metadata(
+    state_dict: Optional[dict],
+    architecture: str,
+    timestamp: float,
+    title: Optional[str] = None,
+    reso: Optional[Union[str, int, Tuple[int, int]]] = None,
+    author: Optional[str] = None,
+    description: Optional[str] = None,
+    license: Optional[str] = None,
+    tags: Optional[str] = None,
+    merged_from: Optional[str] = None,
+    timesteps: Optional[Tuple[int, int]] = None,
+    is_lora: bool = True,
+    custom_arch: Optional[str] = None,
+):
+    metadata = {}
+    metadata.update(BASE_METADATA)
+
+    # TODO implement if we can calculate hash without loading all tensors
+    # if state_dict is not None:
+    # hash = precalculate_safetensors_hashes(state_dict)
+    # metadata["modelspec.hash_sha256"] = hash
+
+    # arch = ARCH_HUNYUAN_VIDEO
+    if architecture == ARCHITECTURE_HUNYUAN_VIDEO:
+        arch = ARCH_HUNYUAN_VIDEO
+        impl = IMPL_HUNYUAN_VIDEO
+    elif architecture == ARCHITECTURE_WAN:
+        arch = ARCH_WAN
+        impl = IMPL_WAN
+    elif architecture == ARCHITECTURE_FRAMEPACK:
+        arch = ARCH_FRAMEPACK
+        impl = IMPL_FRAMEPACK
+    elif architecture == ARCHITECTURE_FLUX_KONTEXT:
+        arch = ARCH_FLUX_KONTEXT
+        impl = IMPL_FLUX_KONTEXT
+    elif (
+        architecture == ARCHITECTURE_FLUX_2_DEV
+        or architecture == ARCHITECTURE_FLUX_2_KLEIN_4B
+        or architecture == ARCHITECTURE_FLUX_2_KLEIN_9B
+    ):
+        if architecture == ARCHITECTURE_FLUX_2_DEV:
+            arch = ARCH_FLUX_2_DEV
+        elif architecture == ARCHITECTURE_FLUX_2_KLEIN_4B:
+            arch = ARCH_FLUX_2_KLEIN_4B
+        elif architecture == ARCHITECTURE_FLUX_2_KLEIN_9B:
+            arch = ARCH_FLUX_2_KLEIN_9B
+        impl = IMPL_FLUX_2
+    elif architecture == ARCHITECTURE_QWEN_IMAGE:
+        arch = ARCH_QWEN_IMAGE
+        impl = IMPL_QWEN_IMAGE
+    elif architecture == ARCHITECTURE_QWEN_IMAGE_EDIT:
+        # We treat Qwen-Image-Edit and Qwen-Image-Edit-Plus the same for architecture and implementation
+        # So we must distinguish them by custom_arch if needed
+        impl = IMPL_QWEN_IMAGE_EDIT
+        if custom_arch is None:
+            arch = ARCH_QWEN_IMAGE_EDIT
+        elif custom_arch == CUSTOM_ARCH_QWEN_IMAGE_EDIT_PLUS:
+            arch = ARCH_QWEN_IMAGE_EDIT_PLUS
+            custom_arch = None  # clear custom_arch to avoid override later
+        elif custom_arch == CUSTOM_ARCH_QWEN_IMAGE_EDIT_2511:
+            arch = ARCH_QWEN_IMAGE_EDIT_2511
+            custom_arch = None  # clear custom_arch to avoid override later
+        else:
+            arch = ARCH_QWEN_IMAGE_EDIT  # override with custom_arch later
+    elif architecture == ARCHITECTURE_QWEN_IMAGE_LAYERED:
+        arch = ARCH_QWEN_IMAGE_LAYERED
+        impl = IMPL_QWEN_IMAGE_LAYERED
+    elif architecture == ARCHITECTURE_KANDINSKY5:
+        arch = ARCH_KANDINSKY5
+        impl = IMPL_KANDINSKY5
+    elif architecture == ARCHITECTURE_HUNYUAN_VIDEO_1_5:
+        arch = ARCH_HUNYUAN_VIDEO_1_5
+        impl = IMPL_HUNYUAN_VIDEO_1_5
+    elif architecture == ARCHITECTURE_Z_IMAGE:
+        arch = ARCH_Z_IMAGE
+        impl = IMPL_Z_IMAGE
+    elif architecture == ARCHITECTURE_HIDREAM_O1:
+        arch = ARCH_HIDREAM_O1
+        impl = IMPL_HIDREAM_O1
+    elif architecture == ARCHITECTURE_IDEOGRAM4:
+        arch = ARCH_IDEOGRAM4
+        impl = IMPL_IDEOGRAM4
+    elif architecture == ARCHITECTURE_KREA2:
+        arch = ARCH_KREA2
+        impl = IMPL_KREA2
+    else:
+        raise ValueError(f"Unknown architecture: {architecture}")
+
+    # Override with custom architecture if provided
+    if custom_arch is not None:
+        arch = custom_arch
+
+    if is_lora:
+        arch += f"/{ADAPTER_LORA}"
+    metadata["modelspec.architecture"] = arch
+
+    metadata["modelspec.implementation"] = impl
+
+    if title is None:
+        title = "LoRA" if is_lora else "Hunyuan-Video"
+        title += f"@{timestamp}"
+    metadata[MODELSPEC_TITLE] = title
+
+    if author is not None:
+        metadata["modelspec.author"] = author
+    else:
+        del metadata["modelspec.author"]
+
+    if description is not None:
+        metadata["modelspec.description"] = description
+    else:
+        del metadata["modelspec.description"]
+
+    if merged_from is not None:
+        metadata["modelspec.merged_from"] = merged_from
+    else:
+        del metadata["modelspec.merged_from"]
+
+    if license is not None:
+        metadata["modelspec.license"] = license
+    else:
+        del metadata["modelspec.license"]
+
+    if tags is not None:
+        metadata["modelspec.tags"] = tags
+    else:
+        del metadata["modelspec.tags"]
+
+    # remove microsecond from time
+    int_ts = int(timestamp)
+
+    # time to iso-8601 compliant date
+    date = datetime.datetime.fromtimestamp(int_ts).isoformat()
+    metadata["modelspec.date"] = date
+
+    if reso is not None:
+        # comma separated to tuple
+        if isinstance(reso, str):
+            reso = tuple(map(int, reso.split(",")))
+        if len(reso) == 1:
+            reso = (reso[0], reso[0])
+    else:
+        # resolution is defined in dataset, so use default here
+        # Use 1328x1328 for Qwen-Image, 1024x1024 for Qwen-Image-Edit and Z-Image, or 1280x720 for others (this is just a placeholder, actual resolution may vary)
+        if architecture == ARCHITECTURE_QWEN_IMAGE:
+            reso = (1328, 1328)
+        elif architecture == ARCHITECTURE_QWEN_IMAGE_EDIT:
+            reso = (1024, 1024)
+        elif architecture == ARCHITECTURE_Z_IMAGE:
+            reso = (1024, 1024)
+        elif architecture == ARCHITECTURE_HIDREAM_O1:
+            reso = (2048, 2048)
+        elif architecture == ARCHITECTURE_IDEOGRAM4:
+            reso = (1024, 1024)
+        elif architecture == ARCHITECTURE_KREA2:
+            reso = (1024, 1024)
+        else:
+            reso = (1280, 720)
+    if isinstance(reso, int):
+        reso = (reso, reso)
+
+    metadata["modelspec.resolution"] = f"{reso[0]}x{reso[1]}"
+
+    # metadata["modelspec.prediction_type"] = PRED_TYPE_EPSILON
+    del metadata["modelspec.prediction_type"]
+
+    if timesteps is not None:
+        if isinstance(timesteps, str) or isinstance(timesteps, int):
+            timesteps = (timesteps, timesteps)
+        if len(timesteps) == 1:
+            timesteps = (timesteps[0], timesteps[0])
+        metadata["modelspec.timestep_range"] = f"{timesteps[0]},{timesteps[1]}"
+    else:
+        del metadata["modelspec.timestep_range"]
+
+    # if clip_skip is not None:
+    #     metadata["modelspec.encoder_layer"] = f"{clip_skip}"
+    # else:
+    del metadata["modelspec.encoder_layer"]
+
+    # # assert all values are filled
+    # assert all([v is not None for v in metadata.values()]), metadata
+    if not all([v is not None for v in metadata.values()]):
+        logger.error(f"Internal error: some metadata values are None: {metadata}")
+
+    return metadata
+
+
+# region utils
+
+
+def get_title(metadata: dict) -> Optional[str]:
+    return metadata.get(MODELSPEC_TITLE, None)
+
+
+def load_metadata_from_safetensors(model: str) -> dict:
+    if not model.endswith(".safetensors"):
+        return {}
+
+    with safetensors.safe_open(model, framework="pt") as f:
+        metadata = f.metadata()
+    if metadata is None:
+        metadata = {}
+    return metadata
+
+
+def build_merged_from(models: List[str]) -> str:
+    def get_title(model: str):
+        metadata = load_metadata_from_safetensors(model)
+        title = metadata.get(MODELSPEC_TITLE, None)
+        if title is None:
+            title = os.path.splitext(os.path.basename(model))[0]  # use filename
+        return title
+
+    titles = [get_title(model) for model in models]
+    return ", ".join(titles)
+
+
+# endregion
+
+
+r"""
+if __name__ == "__main__":
+    import argparse
+    import torch
+    from safetensors.torch import load_file
+    from library import train_util
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt", type=str, required=True)
+    args = parser.parse_args()
+
+    print(f"Loading {args.ckpt}")
+    state_dict = load_file(args.ckpt)
+
+    print(f"Calculating metadata")
+    metadata = get(state_dict, False, False, False, False, "sgm", False, False, "title", "date", 256, 1000, 0)
+    print(metadata)
+    del state_dict
+
+    # by reference implementation
+    with open(args.ckpt, mode="rb") as file_data:
+        file_hash = hashlib.sha256()
+        head_len = struct.unpack("Q", file_data.read(8))  # int64 header length prefix
+        header = json.loads(file_data.read(head_len[0]))  # header itself, json string
+        content = (
+            file_data.read()
+        )  # All other content is tightly packed tensors. Copy to RAM for simplicity, but you can avoid this read with a more careful FS-dependent impl.
+        file_hash.update(content)
+        # ===== Update the hash for modelspec =====
+        by_ref = f"0x{file_hash.hexdigest()}"
+    print(by_ref)
+    print("is same?", by_ref == metadata["modelspec.hash_sha256"])
+
+"""
diff --git a/musubi_tuner/utils/train_utils.py b/musubi_tuner/utils/train_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..82a4b5a7745838a315e5a0c8f97d324b1e2552eb
--- /dev/null
+++ b/musubi_tuner/utils/train_utils.py
@@ -0,0 +1,201 @@
+import argparse
+import logging
+import os
+import shutil
+from typing import Callable
+
+import accelerate
+import torch
+
+from musubi_tuner.utils import huggingface_utils
+from musubi_tuner.utils.model_utils import str_to_dtype
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+# checkpointファイル名
+EPOCH_STATE_NAME = "{}-{:06d}-state"
+EPOCH_FILE_NAME = "{}-{:06d}"
+EPOCH_DIFFUSERS_DIR_NAME = "{}-{:06d}"
+LAST_STATE_NAME = "{}-state"
+STEP_STATE_NAME = "{}-step{:08d}-state"
+STEP_FILE_NAME = "{}-step{:08d}"
+STEP_DIFFUSERS_DIR_NAME = "{}-step{:08d}"
+
+
+def get_sanitized_config_or_none(args: argparse.Namespace):
+    # if `--log_config` is enabled, return args for logging. if not, return None.
+    # when `--log_config is enabled, filter out sensitive values from args
+    # if wandb is not enabled, the log is not exposed to the public, but it is fine to filter out sensitive values to be safe
+
+    if not args.log_config:
+        return None
+
+    sensitive_args = ["wandb_api_key", "huggingface_token"]
+    sensitive_path_args = [
+        "dit",
+        "vae",
+        "text_encoder1",
+        "text_encoder2",
+        "image_encoder",
+        "base_weights",
+        "network_weights",
+        "output_dir",
+        "logging_dir",
+    ]
+    filtered_args = {}
+    for k, v in vars(args).items():
+        # filter out sensitive values and convert to string if necessary
+        if k not in sensitive_args + sensitive_path_args:
+            # Accelerate values need to have type `bool`,`str`, `float`, `int`, or `None`.
+            if v is None or isinstance(v, bool) or isinstance(v, str) or isinstance(v, float) or isinstance(v, int):
+                filtered_args[k] = v
+            # accelerate does not support lists
+            elif isinstance(v, list):
+                filtered_args[k] = f"{v}"
+            # accelerate does not support objects
+            elif isinstance(v, object):
+                filtered_args[k] = f"{v}"
+
+    return filtered_args
+
+
+class LossRecorder:
+    def __init__(self):
+        self.loss_list: list[float] = []
+        self.loss_total: float = 0.0
+
+    def add(self, *, epoch: int, step: int, loss: float) -> None:
+        if epoch == 0:
+            self.loss_list.append(loss)
+        else:
+            while len(self.loss_list) <= step:
+                self.loss_list.append(0.0)
+            self.loss_total -= self.loss_list[step]
+            self.loss_list[step] = loss
+        self.loss_total += loss
+
+    @property
+    def moving_average(self) -> float:
+        return self.loss_total / len(self.loss_list)
+
+
+def get_epoch_ckpt_name(model_name, epoch_no: int):
+    return EPOCH_FILE_NAME.format(model_name, epoch_no) + ".safetensors"
+
+
+def get_step_ckpt_name(model_name, step_no: int):
+    return STEP_FILE_NAME.format(model_name, step_no) + ".safetensors"
+
+
+def get_last_ckpt_name(model_name):
+    return model_name + ".safetensors"
+
+
+def get_remove_epoch_no(args: argparse.Namespace, epoch_no: int):
+    if args.save_last_n_epochs is None:
+        return None
+
+    remove_epoch_no = epoch_no - args.save_every_n_epochs * args.save_last_n_epochs
+    if remove_epoch_no < 0:
+        return None
+    return remove_epoch_no
+
+
+def get_remove_step_no(args: argparse.Namespace, step_no: int):
+    if args.save_last_n_steps is None:
+        return None
+
+    # calculate the step number to remove from the last_n_steps and save_every_n_steps
+    # e.g. if save_every_n_steps=10, save_last_n_steps=30, at step 50, keep 30 steps and remove step 10
+    remove_step_no = step_no - args.save_last_n_steps - 1
+    remove_step_no = remove_step_no - (remove_step_no % args.save_every_n_steps)
+    if remove_step_no < 0:
+        return None
+    return remove_step_no
+
+
+def save_and_remove_state_on_epoch_end(args: argparse.Namespace, accelerator: accelerate.Accelerator, epoch_no: int):
+    model_name = args.output_name
+
+    logger.info("")
+    logger.info(f"saving state at epoch {epoch_no}")
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    state_dir = os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, epoch_no))
+    accelerator.save_state(state_dir)
+    if args.save_state_to_huggingface:
+        logger.info("uploading state to huggingface.")
+        huggingface_utils.upload(args, state_dir, "/" + EPOCH_STATE_NAME.format(model_name, epoch_no))
+
+    last_n_epochs = args.save_last_n_epochs_state if args.save_last_n_epochs_state else args.save_last_n_epochs
+    if last_n_epochs is not None:
+        remove_epoch_no = epoch_no - args.save_every_n_epochs * last_n_epochs
+        state_dir_old = os.path.join(args.output_dir, EPOCH_STATE_NAME.format(model_name, remove_epoch_no))
+        if os.path.exists(state_dir_old):
+            logger.info(f"removing old state: {state_dir_old}")
+            shutil.rmtree(state_dir_old)
+
+
+def save_and_remove_state_stepwise(args: argparse.Namespace, accelerator: accelerate.Accelerator, step_no: int):
+    model_name = args.output_name
+
+    logger.info("")
+    logger.info(f"saving state at step {step_no}")
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    state_dir = os.path.join(args.output_dir, STEP_STATE_NAME.format(model_name, step_no))
+    accelerator.save_state(state_dir)
+    if args.save_state_to_huggingface:
+        logger.info("uploading state to huggingface.")
+        huggingface_utils.upload(args, state_dir, "/" + STEP_STATE_NAME.format(model_name, step_no))
+
+    last_n_steps = args.save_last_n_steps_state if args.save_last_n_steps_state else args.save_last_n_steps
+    if last_n_steps is not None:
+        # last_n_steps前のstep_noから、save_every_n_stepsの倍数のstep_noを計算して削除する
+        remove_step_no = step_no - last_n_steps - 1
+        remove_step_no = remove_step_no - (remove_step_no % args.save_every_n_steps)
+
+        if remove_step_no > 0:
+            state_dir_old = os.path.join(args.output_dir, STEP_STATE_NAME.format(model_name, remove_step_no))
+            if os.path.exists(state_dir_old):
+                logger.info(f"removing old state: {state_dir_old}")
+                shutil.rmtree(state_dir_old)
+
+
+def save_state_on_train_end(args: argparse.Namespace, accelerator: accelerate.Accelerator):
+    model_name = args.output_name
+
+    logger.info("")
+    logger.info("saving last state.")
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    state_dir = os.path.join(args.output_dir, LAST_STATE_NAME.format(model_name))
+    accelerator.save_state(state_dir)
+
+    if args.save_state_to_huggingface:
+        logger.info("uploading last state to huggingface.")
+        huggingface_utils.upload(args, state_dir, "/" + LAST_STATE_NAME.format(model_name))
+
+
+def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15) -> Callable[[float], float]:
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def resolve_save_dtype(save_precision: str | None, full_fp16: bool = False, full_bf16: bool = False) -> torch.dtype:
+    """Resolve the dtype for saving network weights.
+
+    Explicit --save_precision wins; otherwise follow full_fp16/full_bf16 so the
+    saved weights match the training precision; otherwise fp32, the precision
+    the network weights are actually trained in.
+    """
+    if save_precision is not None:
+        return str_to_dtype(save_precision)
+    if full_fp16:
+        return torch.float16
+    if full_bf16:
+        return torch.bfloat16
+    return torch.float32
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..57380774d2a6005b15d483c1e619aeeabb2c06f9
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,10 @@
+torch==2.9.1
+torchvision==0.24.1
+transformers==4.57.6
+accelerate==1.6.0
+huggingface-hub==0.34.3
+safetensors
+av
+numpy
+pillow
+einops
diff --git a/rtcwm/__init__.py b/rtcwm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d537b574a3a9dc69cda9b58ede341e05ebf02e2d
--- /dev/null
+++ b/rtcwm/__init__.py
@@ -0,0 +1,3 @@
+"""Real-time Code World Model: CARLA data generation and causal fine-tuning of the released CWM."""
+
+__version__ = "0.1.0"
diff --git a/rtcwm/data/__init__.py b/rtcwm/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..88017807509f22459f56316b4e37bc4d1796e768
--- /dev/null
+++ b/rtcwm/data/__init__.py
@@ -0,0 +1,2 @@
+"""Data generation: CARLA renders, coarse proxies, captions, QC. Runs in the `carla` environment
+(rendering) or any environment with numpy / PIL / av (proxies, QC); captions need the `cwm` one."""
diff --git a/rtcwm/data/anchors.py b/rtcwm/data/anchors.py
new file mode 100644
index 0000000000000000000000000000000000000000..977146c3baff9705cdfd94d083d75f5ae317f444
--- /dev/null
+++ b/rtcwm/data/anchors.py
@@ -0,0 +1,60 @@
+"""Start frames (Picture 1) for streaming generation.
+
+A stream starts from a start frame: Picture 1, its static caption and the proxy of that frame. Training
+draws Picture 1 from many points of a clip -- the first frame of every ANCHOR_EVERY-th chunk -- either
+right at the start of the sampled span (generation from that image, no history) or earlier in the clip.
+Start frames inside a caption window flagged by QC are skipped. Per start chunk m:
+    CLIP/prompts/static_cMM.txt                      the caption (rtcwm.data.captions --static)
+    CACHE/CLIP/anchor_cMM.png                        the frame
+    CACHE/CLIP/anchor[n]_cMM.safetensors             its Picture 1 latent (n: native 1344x768, else 2048 short edge)
+    CACHE/CLIP/static[n]_cMM.text.safetensors        its Qwen states (rtcwm.train.preencode --stage static_text)
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from cwm_h3_inference.constants import STRIDE_FRAMES, WINDOW_FRAMES
+
+from .. import paths
+from .clip import read_meta
+
+ANCHOR_EVERY = 3                    # chunks between start frames
+CLIP_FRAMES = 17                    # frames per chunk (one VAE clip)
+STREAM_CHUNKS = 54                  # chunks of a 934-frame render's stream
+LAST_SPAN_START = STREAM_CHUNKS - 6
+
+
+def anchor_frame(m: int) -> int:
+    return CLIP_FRAMES * m
+
+
+def _flagged(clip: str):
+    from ..data.qc import lists
+    _, bad = lists.load()
+    return [(STRIDE_FRAMES * w, STRIDE_FRAMES * w + WINDOW_FRAMES) for w in bad.get(clip, ())]
+
+
+def clip_anchors(clip: Path | str) -> list[int]:
+    """Start chunks of a clip: every ANCHOR_EVERY-th chunk up to the last span start, outside flagged windows."""
+    name = clip if isinstance(clip, str) else clip.name
+    flagged = _flagged(name)
+    return [m for m in range(0, LAST_SPAN_START + 1, ANCHOR_EVERY)
+            if not any(a <= anchor_frame(m) < b for a, b in flagged)]
+
+
+def caption_path(clip: Path, m: int) -> Path:
+    return clip / "prompts" / f"static_c{m:02d}.txt"
+
+
+# Picture 1 resolution. "native": the renders' 1344x768 frame as is (1008 H3 rows, ~2.6k Qwen rows; equal or
+# better than the released 2048 short edge, which only upsamples our frames -- ops/dbg/res_test.py).
+# "2048": the released CWM anchor size (7168 H3 rows, ~8.8k Qwen rows), used by stream_df_v1 / v2.
+PICTURE1_RES = ("native", "2048")
+
+
+def cache_paths(clip: str, m: int, res: str = "native") -> dict:
+    assert res in PICTURE1_RES, res
+    d = paths.CACHE / clip
+    tag = "n" if res == "native" else ""
+    return dict(png=d / f"anchor_c{m:02d}.png", latent=d / f"anchor{tag}_c{m:02d}.safetensors",
+                text=d / f"static{tag}_c{m:02d}.text.safetensors")
diff --git a/rtcwm/data/captions.py b/rtcwm/data/captions.py
new file mode 100644
index 0000000000000000000000000000000000000000..828ee8d21e0e7bcdbefff3789a86294f93a93d00
--- /dev/null
+++ b/rtcwm/data/captions.py
@@ -0,0 +1,184 @@
+"""CWM-style captions with Qwen3-VL-32B: per window (the released CWM format) or one static caption per clip.
+
+Each 124-frame window (stride 90) gets 8 evenly spaced frames, one released CWM prompt as the format
+reference, and the ground-truth motion of the protagonist per 0.75 s slot (from the simulated
+trajectory, so a stop or a turn-round hard to see in 8 frames is not missed). Windows after the first
+get CWM's rolling-history preamble. Output: CLIP/prompts/NNN.txt; existing prompts are kept.
+
+    python -m rtcwm.data.captions CLIP_DIR [CLIP_DIR ...]      (env MAX_WINDOWS limits windows per clip)
+    python -m rtcwm.data.captions --static CLIP_DIR [...]   -> CLIP/prompts/static_cMM.txt
+
+Static captions are for streaming generation: a stream starts from a start frame (Picture 1), and the
+caption describes only that frame -- the protagonist and the scene visible in it -- plus one fixed sentence
+on the walking and the trailing camera, so the text carries no timing and nothing from later frames;
+motion comes from the proxy alone. Start frames: the first frame of every ANCHOR_EVERY-th chunk
+(rtcwm.data.anchors), so training draws Picture 1 from many points of a clip.
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+import re
+import sys
+from pathlib import Path
+
+import av
+import numpy as np
+
+from .. import paths
+from .clip import FPS, STRIDE_FRAMES, WINDOW_FRAMES, n_windows, read_meta
+
+REFERENCE_EXAMPLE = "example_30"
+HISTORY = ("The first 1.4167 seconds are the immutable 34-frame rolling history and remain unchanged, continuing the "
+           "established camera motion, positions, scale, orientation, materials, lighting, and ongoing actions without a cut, restart, or reset.\n\n")
+SLOTS = [(0.0, 0.75), (0.75, 1.5), (1.5, 2.25), (2.25, 3.0), (3.0, 3.75), (3.75, 4.5), (4.5, 5.17)]
+N_FRAMES = 8
+
+INSTRUCTION = """You write video captions for a world model. Below is a REFERENCE caption that shows the required format and level of detail.
+
+REFERENCE:
+<<<
+{ref}
+>>>
+
+Now look at the {n} frames, which are evenly spaced over a 5.17-second clip (frame k is at {times}).
+Write a caption for THIS clip in exactly the same structure as the reference:
+1. "Appearance and scene:" then one paragraph describing the protagonist (the person the camera follows, seen mostly from behind): build, hair, clothing, stable identity features, and where they are in frame at 0.00s.
+2. One paragraph describing the setting: ground surface, road, buildings, vegetation, vehicles, sky, lighting and weather, and the camera (a close trailing third-person camera behind the protagonist) and how it moves.
+3. "Protagonist action timeline:" then exactly seven segments "0.00–0.75s: ...; 0.75–1.50s: ...; 1.50–2.25s: ...; 2.25–3.00s: ...; 3.00–3.75s: ...; 3.75–4.50s: ...; 4.50–5.17s: ..." describing walking, turning left or right, slowing, crossing, etc., based on what changes between frames.
+4. "Other visible actors:" then a short paragraph about other cars or people that move, with time ranges, or state that none move.
+Ground-truth motion of the protagonist in this clip (follow it exactly in the action timeline; the trailing camera always swings round to stay behind the protagonist, so after turning around the protagonist is again seen from behind):
+{hints}
+Describe only what is visible. Do not mention frames, images, renders, simulation, CARLA, or games. Output the caption only."""
+
+
+def reference_prompt() -> str:
+    cfg = json.loads((paths.CWM_EXAMPLES / REFERENCE_EXAMPLE / "config.json").read_text())
+    return cfg["windows"][1]["prompt"].replace("\r\n", "\n")
+
+
+def motion_hints(traj, w: int) -> str:
+    t = np.asarray(traj)[:, :2]
+    head = []
+    for i in range(len(t)):
+        d = t[min(len(t) - 1, i + 6)] - t[max(0, i - 6)]
+        head.append(math.degrees(math.atan2(d[1], d[0])) if np.hypot(*d) > 0.3 else None)
+    out = []
+    for s0, s1 in SLOTS:
+        f0, f1 = STRIDE_FRAMES * w + int(s0 * FPS), min(STRIDE_FRAMES * w + int(s1 * FPS), len(t) - 1)
+        dist = float(np.linalg.norm(t[f1] - t[f0]))
+        h0, h1 = head[f0], head[f1]
+        turn = None if h0 is None or h1 is None else abs((h1 - h0 + 180) % 360 - 180)
+        if dist < 0.25:
+            state = "stands still / pauses"
+        elif turn is not None and turn > 100:
+            state = "turns around and starts walking back the way they came"
+        elif turn is not None and turn > 35:
+            state = "walks while turning to follow the path"
+        else:
+            state = "walks steadily forward along the path"
+        out.append(f"{s0:.2f}-{s1:.2f}s: {state}")
+    return "; ".join(out)
+
+
+def clean_caption(text: str) -> str:
+    """Strip prompt delimiters the model sometimes echoes; put the seven timeline segments on one
+    line joined by '; ' as in the released CWM captions."""
+    text = text.replace("<<<", "").replace(">>>", "").strip()
+    if "Protagonist action timeline:" in text:
+        head, rest = text.split("Protagonist action timeline:", 1)
+        tl, sep, tail = rest.partition("Other visible actors:")
+        segs = [re.sub(r"\s+", " ", x).strip().rstrip(";") for x in re.split(r"\n\s*\n|;\s*(?=\d+\.\d+[–-])", tl) if x.strip()]
+        text = head.rstrip() + "\n\nProtagonist action timeline:\n\n" + "; ".join(segs)
+        if sep:
+            text += "\n\nOther visible actors:\n\n" + re.sub(r"\s+", " ", tail).strip()
+    return text
+
+
+STATIC_INSTRUCTION = """You write the text prompt for a world model that continues a video from this single starting image.
+Write exactly two paragraphs, starting with the line "Appearance and scene:":
+1. The protagonist (the person in the foreground, seen from behind): build, hair, skin tone, clothing, shoes and stable
+   identity features, and where they are in the frame.
+2. The setting visible in this image: ground surface, sidewalk, road, buildings, walls, vegetation, vehicles, other
+   people, sky, lighting and weather, and the camera (a close trailing third-person camera behind the protagonist).
+Describe only what is visible in this image. Do not describe actions, events or anything that happens later. Do not
+mention images, renders, simulation, CARLA, or games. Output the two paragraphs only."""
+STATIC_MOTION = ("The protagonist walks along the sidewalk, and the close trailing third-person camera follows behind "
+                 "them, swinging round to stay behind the protagonist whenever they turn.")
+
+
+class Captioner:
+    def __init__(self):
+        import torch
+        from transformers import AutoModelForImageTextToText, AutoProcessor
+        self.torch = torch
+        self.proc = AutoProcessor.from_pretrained(str(paths.qwen_vl()))
+        self.model = AutoModelForImageTextToText.from_pretrained(str(paths.qwen_vl()), torch_dtype=torch.bfloat16,
+                                                                 device_map="auto").eval()
+        self.ref = reference_prompt()
+
+    def caption(self, images, times: str, hints: str) -> str:
+        content = [{"type": "image", "image": im} for im in images]
+        content.append({"type": "text", "text": INSTRUCTION.format(ref=self.ref, n=len(images), times=times, hints=hints)})
+        inputs = self.proc.apply_chat_template([{"role": "user", "content": content}], tokenize=True,
+                                               add_generation_prompt=True, return_dict=True,
+                                               return_tensors="pt").to(self.model.device)
+        with self.torch.inference_mode():
+            out = self.model.generate(**inputs, max_new_tokens=1200, do_sample=False)
+        return self.proc.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0].strip()
+
+    def caption_static(self, clip: Path):
+        from .anchors import anchor_frame, caption_path, clip_anchors
+        todo = [m for m in clip_anchors(clip) if not caption_path(clip, m).exists()]
+        if not todo:
+            return
+        with av.open(str(clip / "rgb.mp4")) as c:
+            frames = {}
+            for i, f in enumerate(c.decode(video=0)):
+                if i in {anchor_frame(m) for m in todo}:
+                    frames[i] = f.to_image()
+        for m in todo:
+            self._caption_frame(frames[anchor_frame(m)], caption_path(clip, m))
+            print(clip.name, "static", m, flush=True)
+
+    def _caption_frame(self, first, outp: Path):
+        content = [{"type": "image", "image": first}, {"type": "text", "text": STATIC_INSTRUCTION}]
+        inputs = self.proc.apply_chat_template([{"role": "user", "content": content}], tokenize=True,
+                                               add_generation_prompt=True, return_dict=True,
+                                               return_tensors="pt").to(self.model.device)
+        with self.torch.inference_mode():
+            out = self.model.generate(**inputs, max_new_tokens=700, do_sample=False)
+        text = self.proc.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0].strip()
+        text = text.replace("<<<", "").replace(">>>", "").strip()
+        outp.parent.mkdir(exist_ok=True)
+        outp.write_text(text + "\n\n" + STATIC_MOTION)
+
+    def caption_clip(self, clip: Path):
+        meta = read_meta(clip)
+        windows = min(n_windows(meta["frames"]), int(os.environ.get("MAX_WINDOWS", 99)))
+        (clip / "prompts").mkdir(exist_ok=True)
+        with av.open(str(clip / "rgb.mp4")) as c:
+            video = [f.to_image() for f in c.decode(video=0)]
+        for w in range(windows):
+            outp = clip / "prompts" / f"{w:03d}.txt"
+            if outp.exists():
+                continue
+            idx = [STRIDE_FRAMES * w + round(k * (WINDOW_FRAMES - 1) / (N_FRAMES - 1)) for k in range(N_FRAMES)]
+            times = ", ".join(f"{k}: {(i - STRIDE_FRAMES * w) / FPS:.2f}s" for k, i in enumerate(idx))
+            text = clean_caption(self.caption([video[i].resize((672, 384)) for i in idx], times,
+                                              motion_hints(meta["hero_traj"], w)))
+            outp.write_text((HISTORY + text) if w > 0 else text)
+            print(clip.name, w, len(text.split()), "words", flush=True)
+
+
+def main():
+    args = sys.argv[1:]
+    static = "--static" in args
+    cap = Captioner()
+    for clip in map(Path, [a for a in args if a != "--static"]):
+        cap.caption_static(clip) if static else cap.caption_clip(clip)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/data/clip.py b/rtcwm/data/clip.py
new file mode 100644
index 0000000000000000000000000000000000000000..86b102acf93a29cb4eb2fbf2471450b99909f9d4
--- /dev/null
+++ b/rtcwm/data/clip.py
@@ -0,0 +1,85 @@
+"""Reading a rendered clip directory.
+
+    rgb.mp4                 1344x768, 24 fps
+    conditions/             fine proxy (CARLA depth + segmentation), NNNNNN.{depth.f32,semantic_id.png}
+    conditions_coarse/      coarse proxy (rtcwm.data.proxy)
+    cameras.json            per-frame UE camera pose + RGB intrinsics
+    scene_static.json       level boxes, lane waypoints, parked vehicles (input of the coarse proxy)
+    scene_frames.jsonl      per-frame protagonist bones and background walkers
+    meta.json               map, seed, route, camera parameters, trajectory, QC
+    prompts/NNN.txt         one caption per 124-frame window
+"""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from .geometry import PROXY_H, PROXY_W
+
+FPS = 24
+WINDOW_FRAMES, STRIDE_FRAMES = 124, 90
+
+
+def n_windows(frames: int) -> int:
+    return (frames - WINDOW_FRAMES) // STRIDE_FRAMES + 1
+
+
+def window_frames(w: int) -> slice:
+    return slice(STRIDE_FRAMES * w, STRIDE_FRAMES * w + WINDOW_FRAMES)
+
+
+def read_meta(clip: Path) -> dict:
+    return json.loads((Path(clip) / "meta.json").read_text())
+
+
+def read_cameras(clip: Path) -> list:
+    return json.loads((Path(clip) / "cameras.json").read_text())["frames"]
+
+
+def read_scene_frames(clip: Path) -> list:
+    with open(Path(clip) / "scene_frames.jsonl") as fh:
+        return [json.loads(line) for line in fh]
+
+
+def read_scene_static(clip: Path) -> dict:
+    return json.loads((Path(clip) / "scene_static.json").read_text())
+
+
+def read_proxy(clip: Path, i: int, kind: str = "conditions"):
+    """(depth [192,336] float32, semantic [192,336] uint8) of frame i."""
+    d = Path(clip) / kind
+    depth = np.fromfile(d / f"{i:06d}.depth.f32", dtype="<f4").reshape(PROXY_H, PROXY_W)
+    sem = np.array(Image.open(d / f"{i:06d}.semantic_id.png"))
+    return depth, sem
+
+
+def qc_pass(clip: Path) -> bool:
+    try:
+        return bool(read_meta(clip).get("qc", {}).get("pass"))
+    except (OSError, ValueError):
+        return False
+
+
+def has_coarse_proxy(clip: Path) -> bool:
+    d = Path(clip) / "conditions_coarse"
+    try:
+        return d.is_dir() and len(os.listdir(d)) >= 2 * read_meta(clip)["frames"]
+    except (OSError, ValueError, KeyError):
+        return False
+
+
+def has_prompts(clip: Path) -> bool:
+    try:
+        nw = n_windows(read_meta(clip)["frames"])
+    except (OSError, ValueError, KeyError):
+        return False
+    return all((Path(clip) / "prompts" / f"{w:03d}.txt").exists() for w in range(nw))
+
+
+def complete(clip: Path) -> bool:
+    """Rendered, passed QC and has its coarse proxy."""
+    return qc_pass(clip) and has_coarse_proxy(clip)
diff --git a/rtcwm/data/geometry.py b/rtcwm/data/geometry.py
new file mode 100644
index 0000000000000000000000000000000000000000..0f7a92271559c4b62cd7a76a6aa92a453c27c376
--- /dev/null
+++ b/rtcwm/data/geometry.py
@@ -0,0 +1,72 @@
+"""Conventions shared by the renderer, the coarse proxy and the QC scans.
+
+CARLA / UE world: x forward, y right, z up (left-handed), metres; rotations are degrees.
+open-code renderer world: x right, y down, z forward.
+CWM proxy frames: 336x192, focal 206.6 px, semantic ids {0,1,2,4,5,8,9,10}.
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+PROXY_W, PROXY_H = 336, 192
+PROXY_F = 206.6
+PROXY_CX, PROXY_CY = PROXY_W / 2, PROXY_H / 2
+HFOV_DEG = math.degrees(2 * math.atan(PROXY_CX / PROXY_F))
+
+# CWM semantic ids
+CLS_OTHER, CLS_SKY, CLS_WATER, CLS_ROAD, CLS_VEGETATION, CLS_HERO, CLS_VEHICLE = 0, 1, 2, 4, 5, 8, 10
+
+FEET_BONES = ("crl_foot__L", "crl_foot__R", "crl_toe__L", "crl_toe__R")
+
+
+def carla_to_renderer(p) -> np.ndarray:
+    return np.array([p[1], -p[2], p[0]], float)
+
+
+def renderer_to_carla(q) -> list:
+    return [float(q[2]), float(q[0]), float(-q[1])]
+
+
+def renderer_axes(yaw_deg: float) -> np.ndarray:
+    """Renderer-space unit axes (rows: forward, right, up) of a CARLA yaw."""
+    y = math.radians(yaw_deg)
+    return np.stack([carla_to_renderer([math.cos(y), math.sin(y), 0.0]),
+                     carla_to_renderer([-math.sin(y), math.cos(y), 0.0]),
+                     carla_to_renderer([0.0, 0.0, 1.0])])
+
+
+def camera_basis(pitch_deg: float, yaw_deg: float):
+    """UE camera axes (forward, right, up) in world coordinates."""
+    p, y = math.radians(pitch_deg), math.radians(yaw_deg)
+    fwd = np.array([math.cos(p) * math.cos(y), math.cos(p) * math.sin(y), math.sin(p)])
+    right = np.array([-math.sin(y), math.cos(y), 0.0])
+    up = np.cross(fwd, right)            # UE is left-handed: this is +up
+    return fwd, right, up
+
+
+def project(points, cam_loc, cam_rot):
+    """World points [N,3] -> proxy pixel (u, v) and planar depth z for a UE camera pose."""
+    fwd, right, up = camera_basis(*cam_rot[:2])
+    q = np.atleast_2d(np.asarray(points, float)) - np.asarray(cam_loc, float)
+    z = q @ fwd
+    with np.errstate(divide="ignore", invalid="ignore"):
+        u = PROXY_CX + PROXY_F * (q @ right) / z
+        v = PROXY_CY - PROXY_F * (q @ up) / z
+    return u, v, z
+
+
+def box_local(xy, centre, yaw_deg: float):
+    """2D offsets of points from a box centre, in the box's own (u, v) axes."""
+    yaw = math.radians(yaw_deg)
+    d = np.atleast_2d(np.asarray(xy, float))[:, :2] - np.asarray(centre[:2], float)
+    return d[:, 0] * math.cos(yaw) + d[:, 1] * math.sin(yaw), -d[:, 0] * math.sin(yaw) + d[:, 1] * math.cos(yaw)
+
+
+def longest_run(flags) -> int:
+    best = run = 0
+    for f in flags:
+        run = run + 1 if f else 0
+        best = max(best, run)
+    return best
diff --git a/rtcwm/data/manifest.py b/rtcwm/data/manifest.py
new file mode 100644
index 0000000000000000000000000000000000000000..159a7104397d51cf90f4b7c59bd7df8971a1fabe
--- /dev/null
+++ b/rtcwm/data/manifest.py
@@ -0,0 +1,41 @@
+"""Render manifests: one JSON row per clip (name, map, seed, seconds).
+
+    python -m rtcwm.data.manifest --n 400 --seed0 30000 --maps Town01_Opt Town03_Opt ... > manifest.jsonl
+
+dataset_v1 (460 rows, seeds 20000-20459, 38.9 s) cycled through Town01/02/03/04/05/10HD with Town02 rows
+remapped to other maps (no out-and-back routes there); its manifest is data_lists/dataset_v1_manifest.jsonl.
+Town04 has since been dropped as well (navmesh under the terrain).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+
+DEFAULT_MAPS = ["Town01_Opt", "Town03_Opt", "Town05_Opt", "Town10HD_Opt"]
+CLIP_SECONDS = 934 / 24                # 10 CWM windows of 124 frames at stride 90
+
+
+def rows(n: int, seed0: int, maps=DEFAULT_MAPS, seconds: float = CLIP_SECONDS, first: int = 0):
+    for i in range(first, first + n):
+        yield dict(name=f"clip_{i:04d}", map=maps[i % len(maps)], seed=seed0 + i, seconds=seconds)
+
+
+def read(path):
+    with open(path) as fh:
+        return [json.loads(line) for line in fh if line.strip()]
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--n", type=int, required=True)
+    ap.add_argument("--seed0", type=int, required=True)
+    ap.add_argument("--first", type=int, default=0, help="index of the first clip (clip_NNNN numbering)")
+    ap.add_argument("--maps", nargs="+", default=DEFAULT_MAPS)
+    ap.add_argument("--seconds", type=float, default=CLIP_SECONDS)
+    a = ap.parse_args()
+    for r in rows(a.n, a.seed0, a.maps, a.seconds, a.first):
+        print(json.dumps(r))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/data/preview.py b/rtcwm/data/preview.py
new file mode 100644
index 0000000000000000000000000000000000000000..019843c2a29b117b697ed9b2b6b9ab1cb86502cf
--- /dev/null
+++ b/rtcwm/data/preview.py
@@ -0,0 +1,82 @@
+"""Preview media of rendered clips: RGB | proxy | overlay video and a 6-frame strip.
+
+    python -m rtcwm.data.preview OUT_DIR CLIP_DIR [...] [--proxy conditions_coarse]
+Writes OUT_DIR/<parent>_<clip>_sbs.mp4, ..._strip.jpg and OUT_DIR/clips.json (one metadata row per clip).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import av
+import cv2
+import numpy as np
+
+from .clip import FPS, read_meta, read_proxy
+
+PALETTE = np.array([[70, 70, 70], [135, 206, 235], [0, 90, 200], [150, 110, 60], [150, 150, 150], [40, 160, 40],
+                    [200, 120, 80], [230, 200, 0], [230, 60, 60], [255, 0, 255], [0, 220, 220], [255, 255, 255]], np.uint8)
+W, H = 448, 256
+
+
+def proxy_image(clip: Path, i: int, kind: str) -> np.ndarray:
+    z, s = read_proxy(clip, i, kind)
+    dn = 1 - np.log(np.clip(z, 0.3, 256) / 0.3) / np.log(256 / 0.3)
+    dn[z <= 1e-3] = 0
+    return cv2.resize((0.55 * PALETTE[s] + 0.45 * (dn[..., None] * 255)).astype(np.uint8), (W, H),
+                      interpolation=cv2.INTER_NEAREST)
+
+
+def label(im, text):
+    im = im.copy()
+    cv2.putText(im, text, (8, 24), 0, 0.6, (0, 0, 0), 4)
+    cv2.putText(im, text, (8, 24), 0, 0.6, (255, 255, 255), 2)
+    return im
+
+
+def preview(clip: Path, out: Path, kind: str = "conditions") -> dict:
+    meta = read_meta(clip)
+    name = f"{clip.parent.name}_{clip.name}"
+    with av.open(str(clip / "rgb.mp4")) as c:
+        frames = [f.to_ndarray(format="rgb24") for f in c.decode(video=0)]
+    with av.open(str(out / f"{name}_sbs.mp4"), "w") as c:
+        st = c.add_stream("libx264", rate=FPS)
+        st.width, st.height, st.pix_fmt = 3 * W, H, "yuv420p"
+        st.options = {"crf": "26"}
+        for i, fr in enumerate(frames):
+            r = cv2.resize(fr, (W, H), interpolation=cv2.INTER_AREA)
+            p = proxy_image(clip, i, kind)
+            row = np.concatenate([label(r, f"RGB {i / FPS:.1f}s"), label(p, "proxy: depth + semantic"),
+                                  label(cv2.addWeighted(r, 0.5, p, 0.5, 0), "overlay")], 1)
+            for pk in st.encode(av.VideoFrame.from_ndarray(np.ascontiguousarray(row), format="rgb24")):
+                c.mux(pk)
+        for pk in st.encode():
+            c.mux(pk)
+    idx = np.linspace(0, len(frames) - 1, 6).astype(int)
+    strip = np.concatenate([label(cv2.resize(frames[i], (W, H)), f"{i / FPS:.1f}s") for i in idx], 1)
+    cv2.imwrite(str(out / f"{name}_strip.jpg"), strip[..., ::-1], [cv2.IMWRITE_JPEG_QUALITY, 82])
+    traj = np.array(meta["hero_traj"])
+    return dict(name=name, map=meta["map"].replace("_Opt", ""), weather=meta["weather"],
+                seconds=round(meta["frames"] / FPS, 1), frames=meta["frames"],
+                hero=round(meta["hero_visible_frac_median"] * 100, 1), fps=round(meta["render_fps"], 1),
+                end_to_start_m=round(float(np.linalg.norm(traj[-1][:2] - traj[0][:2])), 1),
+                path_m=round(float(np.linalg.norm(np.diff(traj[:, :2], axis=0), axis=1).sum()), 1))
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("out", type=Path)
+    ap.add_argument("clips", nargs="+", type=Path)
+    ap.add_argument("--proxy", default="conditions", help="conditions (fine) or conditions_coarse")
+    a = ap.parse_args()
+    a.out.mkdir(parents=True, exist_ok=True)
+    rows = []
+    for clip in a.clips:
+        rows.append(preview(clip, a.out, a.proxy))
+        print(rows[-1]["name"], "ok", flush=True)
+    (a.out / "clips.json").write_text(json.dumps(rows, indent=1))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/data/proxy.py b/rtcwm/data/proxy.py
new file mode 100644
index 0000000000000000000000000000000000000000..127d4c0894aabccf3719f2186c4225aaf60a172c
--- /dev/null
+++ b/rtcwm/data/proxy.py
@@ -0,0 +1,249 @@
+"""Coarse CWM proxy of a rendered clip, compiled from its scene description.
+
+The fine proxy (conditions/) is CARLA's own depth + segmentation of the real meshes. At test time a
+coding agent only has coarse geometry, which is what CWM's proxy is. This rebuilds that with
+open-code-world-model's primitives and ray tracer, from the exact camera of every RGB frame:
+
+    buildings, walls, props -> their CARLA level boxes, trimmed clear of the walk and camera path
+                               (Town10HD building boxes include awnings over the sidewalk)
+    poles, lights, signs    -> a capsule (+ a small head box); their boxes include the lamp arm
+    vegetation              -> open-code lobed clumps, split along long boxes, lobes kept inside the
+                               box (its Gaussian scatter once hung a lobe over the camera)
+    roads                   -> flat strips from driving-lane waypoints; ground -> one plane
+    backdrop                -> a ring of distant boxes (open-code: 320-900 m), so the skyline is geometry
+    protagonist             -> open-code's generic Human rig (class 8): position, heading and a
+                               procedural gait only -- no body shape or clothing
+    other walkers           -> the same rig (class 0); vehicles -> wireframe boxes (class 10)
+
+    python -m rtcwm.data.proxy [--workers 32] [--out conditions_coarse] CLIP_DIR ...
+"""
+from __future__ import annotations
+
+import argparse
+import math
+from multiprocessing import Pool
+from pathlib import Path
+
+import numpy as np
+
+from .. import paths
+from .clip import FPS, read_cameras, read_meta, read_scene_frames, read_scene_static
+from .geometry import CLS_HERO, CLS_OTHER, CLS_ROAD, CLS_VEGETATION, CLS_VEHICLE, CLS_WATER, PROXY_F
+from .geometry import carla_to_renderer as R
+from .geometry import renderer_axes as axes
+
+paths.use_ocwm()
+from compiler import writer                                                     # noqa: E402
+from compiler.render import Box, Camera, Capsule, Plane, render_frame, vegetation_clump, wireframe_box  # noqa: E402
+from compiler.rig import Human, idle_pose, walk_pose                            # noqa: E402
+
+CULL_M = 200.0
+HERO_SCALE = 1.75 / 1.60          # open-code rig is 1.60 m; CARLA adults are ~1.75 m
+WALKER_RADIUS_M = 60.0
+PATH_CLEARANCE_M = 0.4
+
+
+def covers_road(b, lanes_xy) -> bool:
+    """A box whose footprint contains a driving-lane waypoint: whole-block building meshes have
+    boxes that span the street, and a coarse box there would wall the road off."""
+    c, e, yaw = b["c"], b["e"], math.radians(b["r"][1])
+    d = lanes_xy - np.array(c[:2])
+    u = d[:, 0] * math.cos(yaw) + d[:, 1] * math.sin(yaw)
+    v = -d[:, 0] * math.sin(yaw) + d[:, 1] * math.cos(yaw)
+    return bool(np.any((np.abs(u) < e[0] - 0.5) & (np.abs(v) < e[1] - 0.5)))
+
+
+def clear_of_path(c, e, yaw_deg, pts, margin=PATH_CLEARANCE_M):
+    """Trim a solid box so no path point lies within `margin` of it: each offending point pushes the
+    nearer of the box's in-plane faces back past it. None if nothing is left."""
+    c, e = np.array(c, float), np.array(e, float)
+    yaw = math.radians(yaw_deg)
+    ax_u, ax_v = np.array([math.cos(yaw), math.sin(yaw)]), np.array([-math.sin(yaw), math.cos(yaw)])
+    lo, hi = -e[:2].copy(), e[:2].copy()
+    for p in pts:
+        if not (c[2] - e[2] - 0.5 <= p[2] <= c[2] + e[2] + 0.5):
+            continue
+        d = np.array(p[:2]) - c[:2]
+        q = np.array([d @ ax_u, d @ ax_v])
+        if not np.all((q > lo - margin) & (q < hi + margin)):
+            continue
+        moves = [(hi[k] - (q[k] - margin), k, "hi") for k in range(2)] + \
+                [((q[k] + margin) - lo[k], k, "lo") for k in range(2)]
+        _, k, side = min(moves)                     # cheapest single-face move (box area removed)
+        if side == "hi":
+            hi[k] = q[k] - margin
+        else:
+            lo[k] = q[k] + margin
+        if hi[k] - lo[k] < 0.2:
+            return None
+    mid = (lo + hi) / 2
+    c2, e2 = c.copy(), e.copy()
+    c2[:2] = c[:2] + mid[0] * ax_u + mid[1] * ax_v
+    e2[:2] = (hi - lo) / 2
+    return c2, e2
+
+
+def vegetation_prims(c, e, yaw_deg):
+    """Lobed clumps for one vegetation box; trees (>= 2.4 m tall) get a canopy on a trunk."""
+    yaw = math.radians(yaw_deg)
+    long_x = e[0] >= e[1]
+    major, minor = (e[0], e[1]) if long_x else (e[1], e[0])
+    n = max(1, int(round(major / max(minor, 0.5))))
+    ux = (math.cos(yaw), math.sin(yaw)) if long_x else (-math.sin(yaw), math.cos(yaw))
+    half = max(major / n, minor)
+    top, bot = c[2] + e[2], c[2] - e[2]
+
+    def inside(lobes, zmin):
+        kept = []
+        for s in lobes:
+            p = np.array([s.centre[2], s.centre[0], -s.centre[1]])        # renderer -> CARLA
+            d = p[:2] - np.array(c[:2])
+            u = d[0] * math.cos(yaw) + d[1] * math.sin(yaw)
+            v = -d[0] * math.sin(yaw) + d[1] * math.cos(yaw)
+            if abs(u) <= e[0] and abs(v) <= e[1] and zmin <= p[2] <= top:
+                kept.append(s)
+        return kept
+
+    prims = []
+    for k in range(n):
+        t = -major + (2 * k + 1) * major / n
+        pc = [c[0] + t * ux[0], c[1] + t * ux[1], c[2]]
+        seed = int(abs(pc[0] * 7.3 + pc[1] * 3.1)) % 10000
+        if e[2] < 1.2:
+            prims += inside(vegetation_clump(R(pc), 0.45 * half, aspect=e[2] / max(half, 0.1), seed=seed), bot)
+        else:
+            ch = 0.35 * e[2]
+            cc = [pc[0], pc[1], top - ch]
+            rad = 0.45 * half
+            prims += inside(vegetation_clump(R(cc), rad, aspect=ch / max(rad, 0.1), seed=seed), max(top - 2 * ch, bot + 2.0))
+            prims.append(Capsule(R([pc[0], pc[1], bot]), R(cc), 0.18, CLS_VEGETATION))
+    return prims
+
+
+def static_prims(scene, centre, ground_z, keep_clear=()):
+    prims = []
+    lanes_xy = np.array([[p[0], p[1]] for p in scene["lanes"]]) if scene["lanes"] else np.zeros((0, 2))
+    for b in scene["boxes"]:
+        c, e = b["c"], b["e"]
+        if math.hypot(c[0] - centre[0], c[1] - centre[1]) > CULL_M or max(e) > 80.0:
+            continue                                            # whole-map landscape boxes swallow the camera
+        # coarse means coarse: no grass clumps or small props a coding agent would never place
+        if (b["cls"] == CLS_VEGETATION and e[2] < 0.4) or (b["cls"] == CLS_OTHER and max(e) < 0.25):
+            continue
+        if b["label"] in ("Static", "Dynamic") and max(e) < 1.0:
+            continue
+        if b["label"] in ("Buildings", "Walls") and len(lanes_xy) and covers_road(b, lanes_xy):
+            continue
+        e = [abs(x) for x in e]
+        ax = axes(b["r"][1])
+        if b["label"] in ("Poles", "TrafficLight", "TrafficSigns"):
+            bot, top = c[2] - e[2], c[2] + e[2]
+            prims.append(Capsule(R([c[0], c[1], bot]), R([c[0], c[1], top]), 0.12, CLS_OTHER))
+            if b["label"] != "Poles":
+                prims.append(Box(R([c[0], c[1], top - 0.4]), np.array([0.3, 0.3, 0.4]), ax, class_id=CLS_OTHER))
+        elif b["cls"] == CLS_VEGETATION:
+            prims += vegetation_prims(c, e, b["r"][1])
+        elif b["cls"] == CLS_WATER:
+            prims.append(Box(R(c), np.array([e[0], e[1], max(e[2], 0.02)]), ax, class_id=CLS_WATER))
+        else:
+            if len(keep_clear):
+                cut = clear_of_path(c, e, b["r"][1], keep_clear)
+                if cut is None:
+                    continue
+                c, e = list(cut[0]), list(cut[1])
+            prims.append(Box(R(c), np.array(e), ax, class_id=CLS_OTHER))
+    for x, y, z, yaw, width in scene["lanes"]:
+        if math.hypot(x - centre[0], y - centre[1]) <= CULL_M and abs(z - ground_z) <= 1.5:   # not another level
+            prims.append(Box(R([x, y, ground_z + 0.01]), np.array([1.05, width / 2.0, 0.01]), axes(yaw), class_id=CLS_ROAD))
+    rng = np.random.default_rng(int(abs(centre[0] * 13 + centre[1] * 7)) % 2**31)
+    for k in range(11):                                         # distant backdrop ring
+        ang = 2 * math.pi * k / 11
+        dist = float(rng.uniform(320.0, 900.0))
+        size = [float(rng.uniform(60.0, 160.0)), float(rng.uniform(60.0, 160.0)), float(rng.uniform(45.0, 130.0))]
+        bx = [centre[0] + dist * math.cos(ang), centre[1] + dist * math.sin(ang)]
+        prims.append(Box(R([bx[0], bx[1], ground_z + size[2] / 2]), np.array([s / 2 for s in size]),
+                         axes(math.degrees(ang)), class_id=CLS_OTHER))
+    for v in scene["vehicles"]:
+        if math.hypot(v["c"][0] - centre[0], v["c"][1] - centre[1]) <= CULL_M:
+            prims += wireframe_box(R(v["c"]), np.array(v["e"]), axes(v["r"][1]), class_id=CLS_VEHICLE)
+    return prims
+
+
+def protagonist(x, y, ground_z, yaw, phase, moving):
+    pose = walk_pose(phase) if moving else idle_pose(phase)
+    return Human(position=R([x, y, ground_z]), facing_deg=yaw, pose=pose, scale=HERO_SCALE).primitives()
+
+
+def walkers(poses, ground_z, t):
+    prims = []
+    for k, (x, y, _z, yaw) in enumerate(poses):
+        h = Human(position=R([x, y, ground_z]), facing_deg=yaw, pose=walk_pose((t / 0.95 + 0.37 * k) % 1.0))
+        for p in h.primitives():
+            p.class_id = CLS_OTHER                              # other people are class 0 in CWM
+            prims.append(p)
+    return prims
+
+
+def headings(xy):
+    """Protagonist heading from the smoothed motion (degrees, CARLA yaw)."""
+    yaw, out = None, []
+    for i in range(len(xy)):
+        d = xy[min(len(xy) - 1, i + 4)] - xy[max(0, i - 4)]
+        if np.linalg.norm(d) > 0.15:
+            target = math.degrees(math.atan2(d[1], d[0]))
+            yaw = target if yaw is None else yaw + ((target - yaw + 180) % 360 - 180) * 0.35
+        out.append(yaw if yaw is not None else 0.0)
+    return out
+
+
+def _render(job):
+    i, cam, walkers_i, static, ground_z, hero = job
+    prims = list(static)
+    prims.append(Plane(np.array([0.0, -ground_z, 0.0]), np.array([0.0, -1.0, 0.0]), class_id=CLS_OTHER))
+    prims += protagonist(*hero)
+    near = [w for w in walkers_i if math.hypot(w[0] - cam["loc"][0], w[1] - cam["loc"][1]) < WALKER_RADIUS_M]
+    prims += walkers(near, ground_z, i / FPS)
+    c = Camera(focal_px=PROXY_F, position=R(cam["loc"]), pitch_deg=cam["rot"][0], yaw_deg=cam["rot"][1], roll_deg=0.0)
+    z, ids = render_frame(c, prims)
+    return i, z, ids
+
+
+def compile_clip(clip: Path, out_name: str = "conditions_coarse", workers: int = 32) -> int:
+    clip = Path(clip)
+    scene, cams, frames = read_scene_static(clip), read_cameras(clip), read_scene_frames(clip)
+    traj = np.array(read_meta(clip)["hero_traj"])
+    feet0 = [v[2] for k, v in frames[0]["bones"].items() if "foot" in k or "toe" in k]
+    path_pts = [list(p) for p in traj[::4]] + [c["loc"] for c in cams[::4]]
+    static = static_prims(scene, traj[:, :2].mean(0), (min(feet0) - 0.03) if feet0 else traj[0, 2] - 0.95,
+                          keep_clear=path_pts)
+    xy = traj[:, :2]
+    step = np.r_[0.0, np.linalg.norm(np.diff(xy, axis=0), axis=1)]
+    dist = np.cumsum(step)
+    speed = np.convolve(step, np.ones(9) / 9, mode="same") * FPS
+    yaws = headings(xy)
+    jobs = []
+    for i, (cam, sf) in enumerate(zip(cams, frames)):
+        feet = [v[2] for k, v in sf["bones"].items() if "foot" in k or "toe" in k]
+        ground_z = min(feet) - 0.03 if feet else traj[i, 2] - 0.95
+        moving = speed[i] > 0.25
+        phase = (dist[i] / 1.5) % 1.0 if moving else (i / 48.0) % 1.0     # 1.5 m per stride cycle
+        jobs.append((i, cam, sf["walkers"], static, ground_z, (traj[i, 0], traj[i, 1], ground_z, yaws[i], phase, moving)))
+    with Pool(workers) as pool:
+        for i, z, ids in pool.imap_unordered(_render, jobs, chunksize=4):
+            writer.write_frame(clip / out_name, i, z, ids)
+    return len(static)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("clips", nargs="+", type=Path)
+    ap.add_argument("--out", default="conditions_coarse", help="output directory name inside each clip")
+    ap.add_argument("--workers", type=int, default=32)
+    a = ap.parse_args()
+    for clip in a.clips:
+        n = compile_clip(clip, a.out, a.workers)
+        print(clip.name, "coarse proxy:", n, "static primitives", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/data/qc/__init__.py b/rtcwm/data/qc/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a25489031914d78e26d67e60abb2e759597c77e
--- /dev/null
+++ b/rtcwm/data/qc/__init__.py
@@ -0,0 +1,10 @@
+"""Dataset QC from saved outputs (no re-rendering).
+
+    python -m rtcwm.data.qc clipping CLIP_DIR ...   interpenetration (also run by the renderer)
+    python -m rtcwm.data.qc audit    CLIP_DIR ...   camera, motion, exposure, coarse-vs-fine audit
+    python -m rtcwm.data.qc coarse   CLIP_DIR ...   coarse proxy swallowing the frame / losing the protagonist
+    python -m rtcwm.data.qc lists                   data_lists/bad_windows.json from the reports
+
+Reports go to paths.QC/<check>/<clip>.json; every report has per-window verdicts (window w = frames
+90w .. 90w+123) so training can drop single windows.
+"""
diff --git a/rtcwm/data/qc/__main__.py b/rtcwm/data/qc/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1a0897b71a2c61a91b6e5346b424b6b9ceccaee
--- /dev/null
+++ b/rtcwm/data/qc/__main__.py
@@ -0,0 +1,47 @@
+"""python -m rtcwm.data.qc {clipping,audit,coarse} [--workers N] CLIP_DIR ... | lists"""
+from __future__ import annotations
+
+import argparse
+import json
+from multiprocessing import Pool
+from pathlib import Path
+
+from ... import paths
+from . import audit, clipping, coarse, lists
+
+SCANS = dict(clipping=clipping.scan, audit=audit.audit, coarse=coarse.scan)
+
+
+def _run(job):
+    check, clip = job
+    try:
+        rep = SCANS[check](Path(clip))
+    except Exception as e:                                      # report and keep going
+        return f"{Path(clip).name} {check} ERROR {e!r}"
+    out = paths.QC / check
+    out.mkdir(parents=True, exist_ok=True)
+    (out / f"{Path(clip).name}.json").write_text(json.dumps(rep))
+    return f"{Path(clip).name} {check} bad_windows={rep.get('bad_windows', '-')}"
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("check", choices=[*SCANS, "lists"])
+    ap.add_argument("clips", nargs="*")
+    ap.add_argument("--workers", type=int, default=1)
+    a = ap.parse_intermixed_args()          # options may sit between the check name and the clips
+    if a.check == "lists":
+        print(lists.build())
+        return
+    jobs = [(a.check, c) for c in a.clips]
+    if a.workers > 1:
+        with Pool(a.workers) as p:
+            for line in p.imap_unordered(_run, jobs):
+                print(line, flush=True)
+    else:
+        for j in jobs:
+            print(_run(j), flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/data/qc/audit.py b/rtcwm/data/qc/audit.py
new file mode 100644
index 0000000000000000000000000000000000000000..468b847ada513063901a65533f35aa95436abaa3
--- /dev/null
+++ b/rtcwm/data/qc/audit.py
@@ -0,0 +1,115 @@
+"""Wider render audit (all from saved outputs).
+
+  camera : single-frame pops (a step > 2.5x both neighbours), roll
+  hero   : walking backwards (toe direction vs velocity), standing still, off screen, too small /
+           too large, other walkers or parked cars overlapping the protagonist
+  video  : black / blown-out / duplicated frames, mean luma, exposure changes (largest frame-to-frame
+           luma jump, largest luma swing within 1 s -- CARLA's auto exposure adapting)
+  proxy  : protagonist present in the fine proxy but missing from the coarse one, coarse-vs-fine
+           depth disagreement
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import av
+import numpy as np
+
+from ..clip import FPS, n_windows, read_cameras, read_meta, read_proxy, read_scene_frames, read_scene_static, window_frames
+from ..geometry import CLS_HERO, CLS_SKY, PROXY_H, PROXY_W, longest_run, project
+
+
+def camera_pops(cams) -> np.ndarray:
+    """Frames where the camera jumps: position or yaw step much larger than both neighbours. The
+    smooth ~0.6 s swing round the protagonist at the turn-round is FollowCamera behaviour, not a pop."""
+    C = np.array([c["loc"] for c in cams])
+    rot = np.array([c["rot"] for c in cams])
+    s = np.r_[0, np.linalg.norm(np.diff(C, axis=0), axis=1), 0]
+    y = np.r_[0, np.abs((np.diff(rot[:, 1]) + 180) % 360 - 180), 0]
+    pop = np.zeros(len(C), bool)
+    for i in range(1, len(C)):
+        pop[i] = (s[i] > 0.15 and s[i] > 2.5 * max(s[i - 1], s[i + 1], 0.02)) or \
+                 (y[i] > 4 and y[i] > 2.5 * max(y[i - 1], y[i + 1], 0.3))
+    return pop
+
+
+def video_stats(path: Path):
+    luma, blown, dup, prev = [], [], [], None
+    with av.open(str(path)) as c:
+        s = c.streams.video[0]
+        s.thread_type = "AUTO"
+        for f in c.decode(s):
+            g = f.reformat(width=PROXY_W, height=PROXY_H, format="gray").to_ndarray().astype(np.int16)
+            luma.append(float(g.mean()))
+            blown.append(float((g >= 250).mean()))
+            dup.append(prev is not None and float(np.abs(g - prev).mean()) < 0.05)
+            prev = g
+    return np.array(luma), np.array(blown), np.array(dup)
+
+
+def _touching(p, z, walkers, vehicles) -> bool:
+    for w in walkers:
+        if math.hypot(w[0] - p[0], w[1] - p[1]) < 0.45 and abs(w[2] - z) < 1.0:
+            return True
+    for v in vehicles:
+        yaw = math.radians(v["r"][1])
+        d = p - np.array(v["c"][:2])
+        lx, ly = d[0] * math.cos(yaw) + d[1] * math.sin(yaw), -d[0] * math.sin(yaw) + d[1] * math.cos(yaw)
+        if abs(lx) < v["e"][0] + 0.15 and abs(ly) < v["e"][1] + 0.15:
+            return True
+    return False
+
+
+def audit(clip: Path) -> dict:
+    clip = Path(clip)
+    meta, cams, frames = read_meta(clip), read_cameras(clip), read_scene_frames(clip)
+    vehicles = read_scene_static(clip)["vehicles"]
+    n = min(len(cams), len(frames))
+    traj = np.array(meta["hero_traj"])[:n]
+    rot = np.array([c["rot"] for c in cams[:n]])
+    pops = camera_pops(cams[:n])
+
+    vel = np.zeros((n, 2))
+    vel[6:-6] = (traj[12:, :2] - traj[:-12, :2]) / (12 / FPS)
+    speed = np.linalg.norm(vel, axis=1)
+    back = np.zeros(n, bool)
+    for i in range(n):
+        b = frames[i]["bones"]
+        face = sum(np.array(b[f"crl_toe__{s}"][:2]) - np.array(b[f"crl_foot__{s}"][:2]) for s in "LR")
+        if speed[i] > 0.5 and np.linalg.norm(face) > 1e-3:
+            back[i] = face @ vel[i] / (np.linalg.norm(face) * speed[i]) < -0.3
+    touch = np.array([_touching(traj[i, :2], traj[i, 2], frames[i]["walkers"], vehicles) for i in range(n)])
+
+    hero_frac, coarse_miss, depth_err, offscreen = [], [], [], []
+    for i in range(0, n, 4):
+        df, sf = read_proxy(clip, i)
+        dc, sc = read_proxy(clip, i, "conditions_coarse")
+        hf = sf == CLS_HERO
+        hero_frac.append(float(hf.mean()))
+        coarse_miss.append(hf.sum() > 50 and (hf & (sc == CLS_HERO)).sum() < 0.3 * hf.sum())
+        u, v, z = project(frames[i]["bones"]["crl_hips__C"], cams[i]["loc"], cams[i]["rot"])
+        offscreen.append(not (z[0] > 0 and 0 <= u[0] < PROXY_W and 0 <= v[0] < PROXY_H))
+        bg = (df > 0) & (dc > 0) & (sf != CLS_SKY) & (sc != CLS_SKY)
+        depth_err.append(float(np.median(np.abs(np.log(dc[bg] / df[bg])))) if bg.sum() > 100 else 0.0)
+    hero_frac, depth_err = np.array(hero_frac), np.array(depth_err)
+    luma, blown, dup = video_stats(clip / "rgb.mp4")
+    m = min(len(luma), n)
+
+    windows = []
+    for w in range(n_windows(n)):
+        sl = window_frames(w)
+        a, b = sl.start, sl.stop
+        a4, b4 = a // 4, (b + 3) // 4
+        windows.append(dict(
+            w=w, camera_pops=int(pops[a:b].sum()), max_roll=float(np.abs(rot[a:b, 2]).max()),
+            backwards_s=longest_run(back[a:b]) / FPS, still_s=longest_run(speed[a:b] < 0.2) / FPS,
+            touch_s=longest_run(touch[a:b]) / FPS, offscreen_s=longest_run(offscreen[a4:b4]) * 4 / FPS,
+            hero_frac_min=float(hero_frac[a4:b4].min()), hero_frac_max=float(hero_frac[a4:b4].max()),
+            coarse_hero_miss=int(sum(coarse_miss[a4:b4])), depth_logerr_med=float(np.median(depth_err[a4:b4])),
+            luma=float(luma[a:min(b, m)].mean()), luma_min=float(luma[a:min(b, m)].min()),
+            blown_max=float(blown[a:min(b, m)].max()), dup_frames=int(dup[a:min(b, m)].sum()),
+            luma_jump_max=float(np.abs(np.diff(luma[a:min(b, m)])).max()),
+            luma_swing_1s=float(max(np.ptp(luma[i:i + FPS]) for i in range(a, max(a + 1, min(b, m) - FPS))))))
+    return dict(clip=clip.name, map=meta["map"], weather=meta.get("weather"), frames=n, video_frames=int(len(luma)),
+                windows=windows)
diff --git a/rtcwm/data/qc/clipping.py b/rtcwm/data/qc/clipping.py
new file mode 100644
index 0000000000000000000000000000000000000000..5dbfd55c6662bbf1052c1a7a166ad65af61bfe52
--- /dev/null
+++ b/rtcwm/data/qc/clipping.py
@@ -0,0 +1,129 @@
+"""Interpenetration in a rendered clip, from the fine depth/semantic, camera poses and bones.
+
+  buried   : the protagonist's feet are below the visible ground -- at the projected toe/foot
+             bones the depth buffer shows a surface clearly in front of the bone (Town04).
+  inside   : the protagonist's hips are over the footprint of a hedge, planter, wall, fence or
+             static prop that spans its shins.
+  pierced  : a pole or other thin solid runs through the body: in the middle 60% of the
+             protagonist's box, a column that is mostly non-protagonist surface at its own depth.
+  occluded : something (typically a wall the camera is inside) covers > 30% of the protagonist's
+             projected bone box with surfaces well in front of it and raised above its feet
+             (the ground in front of the feet and walls beside the path do not count).
+
+A window is bad if buried / inside / occluded holds for > 0.5 s in a row, or pierced for >= 3 frames
+(walking through a post takes ~0.3 s).
+"""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+
+from ..clip import FPS, n_windows, read_cameras, read_proxy, read_scene_frames, read_scene_static, window_frames
+from ..geometry import (CLS_HERO, CLS_SKY, FEET_BONES, PROXY_CX, PROXY_CY, PROXY_F, PROXY_H, PROXY_W, camera_basis,
+                        longest_run)
+
+RUN = int(0.5 * FPS)
+PIERCED_RUN = 3
+SOLID_LABELS = ("Walls", "Fences", "GuardRail", "Static")   # + vegetation under 2.4 m (hedges, planters)
+
+
+def _solid_boxes(clip: Path):
+    """Low solid things a walker must not stand in. Tree boxes are excluded (their canopy footprint
+    covers sidewalks), and so are buildings (their boxes include awnings)."""
+    try:
+        boxes = read_scene_static(clip)["boxes"]
+    except OSError:
+        return []
+    out = []
+    for b in boxes:
+        e = [abs(x) for x in b["e"]]
+        if (b["label"] in SOLID_LABELS or (b["label"] == "Vegetation" and e[2] < 1.2)) and max(e[0], e[1]) > 0.2:
+            out.append((np.array(b["c"]), np.array(e), math.radians(b["r"][1])))
+    return out
+
+
+def _buried(d, s, bones, C, fwd, right, up) -> bool:
+    sunk = 0
+    for b in FEET_BONES:
+        q = np.array(bones[b]) - C
+        z = q @ fwd
+        if z < 0.3:
+            continue
+        iu, iv = int(PROXY_CX + PROXY_F * (q @ right) / z), int(PROXY_CY - PROXY_F * (q @ up) / z)
+        if not (1 <= iu < PROXY_W - 1 and 1 <= iv < PROXY_H - 1):
+            continue
+        win_d, win_s = d[iv - 1:iv + 2, iu - 1:iu + 2], s[iv - 1:iv + 2, iu - 1:iu + 2]
+        if not (win_s == CLS_HERO).any() and np.median(win_d) < z - 0.15:
+            sunk += 1
+    return sunk >= 3
+
+
+def _inside(bones, foot_z, solid) -> bool:
+    hp = np.array(bones["crl_hips__C"])
+    for bc, be, yaw in solid:
+        if not (bc[2] - be[2] < foot_z + 0.5 and bc[2] + be[2] > foot_z + 0.3):
+            continue
+        dx, dy = hp[0] - bc[0], hp[1] - bc[1]
+        u, v = dx * math.cos(yaw) + dy * math.sin(yaw), -dx * math.sin(yaw) + dy * math.cos(yaw)
+        if abs(u) < be[0] - 0.1 and abs(v) < be[1] - 0.1:
+            return True
+    return False
+
+
+def _box_checks(d, s, bones, C, fwd, right, up, foot_z, xn, yn):
+    """(occluded fraction of the protagonist's box, pole through the body)."""
+    hd = (np.array(bones["crl_hips__C"]) - C) @ fwd
+    pts = np.array(list(bones.values())) - C
+    zs = pts @ fwd
+    if not (zs > 0.3).all():
+        return 0.0, False
+    us, vs = PROXY_CX + PROXY_F * (pts @ right) / zs, PROXY_CY - PROXY_F * (pts @ up) / zs
+    u0, u1 = int(max(0, us.min() - 3)), int(min(PROXY_W, us.max() + 4))
+    v0, v1 = int(max(0, vs.min() - 3)), int(min(PROXY_H, vs.max() + 4))
+    if u1 - u0 < 3 or v1 - v0 < 3:
+        return 0.0, False
+    bd, bs = d[v0:v1, u0:u1], s[v0:v1, u0:u1]
+    height = C[2] + bd * (fwd[2] + xn[v0:v1, u0:u1] * right[2] + yn[v0:v1, u0:u1] * up[2]) - foot_z
+    occ = float(((bd > 0) & (bd < 0.75 * hd) & (height > 0.3) & (bs != CLS_HERO) & (bs != CLS_SKY)).mean())
+    hz = np.median(bd[bs == CLS_HERO]) if (bs == CLS_HERO).sum() > 20 else hd
+    same = (np.abs(bd - hz) < 0.3) & (bs != CLS_HERO) & (bs != CLS_SKY) & (height > 0.25)
+    w0 = int(0.2 * (u1 - u0))
+    cols = same[:, w0:(u1 - u0) - w0].mean(0) if (u1 - u0) - 2 * w0 > 0 else np.zeros(1)
+    return occ, bool(cols.max() > 0.5)
+
+
+def scan(clip: Path) -> dict:
+    clip = Path(clip)
+    cams, frames = read_cameras(clip), read_scene_frames(clip)
+    n = min(len(cams), len(frames))
+    vv, uu = np.mgrid[0:PROXY_H, 0:PROXY_W]
+    xn, yn = (uu + 0.5 - PROXY_CX) / PROXY_F, -(vv + 0.5 - PROXY_CY) / PROXY_F
+    solid = _solid_boxes(clip)
+    buried, inside, occluded, occ_frac, pierced = [], [], [], [], []
+    for i in range(n):
+        d, s = read_proxy(clip, i)
+        C = np.array(cams[i]["loc"])
+        fwd, right, up = camera_basis(*cams[i]["rot"][:2])
+        bones = frames[i]["bones"]
+        foot_z = min(bones[b][2] for b in FEET_BONES)
+        buried.append(_buried(d, s, bones, C, fwd, right, up))
+        inside.append(_inside(bones, foot_z, solid))
+        frac, pole = _box_checks(d, s, bones, C, fwd, right, up, foot_z, xn, yn)
+        occ_frac.append(round(frac, 3))
+        occluded.append(frac > 0.3)
+        pierced.append(pole)
+    windows = []
+    for w in range(n_windows(n)):
+        sl = window_frames(w)
+        r = dict(w=w, buried_run=longest_run(buried[sl]), occluded_run=longest_run(occluded[sl]),
+                 inside_run=longest_run(inside[sl]), pierced_run=longest_run(pierced[sl]))
+        r["bad"] = r["pierced_run"] >= PIERCED_RUN or max(r["buried_run"], r["occluded_run"], r["inside_run"]) > RUN
+        windows.append(r)
+    return dict(clip=clip.name, frames=n,
+                max_buried_run=longest_run(buried), max_occluded_run=longest_run(occluded),
+                max_inside_run=longest_run(inside), max_pierced_run=longest_run(pierced),
+                bad_windows=[x["w"] for x in windows if x["bad"]], windows=windows,
+                buried=[int(b) for b in buried], inside=[int(b) for b in inside],
+                pierced=[int(b) for b in pierced], occ_frac=occ_frac)
diff --git a/rtcwm/data/qc/coarse.py b/rtcwm/data/qc/coarse.py
new file mode 100644
index 0000000000000000000000000000000000000000..dba7bebd1ed647e344b4f5947516180a0826f5d4
--- /dev/null
+++ b/rtcwm/data/qc/coarse.py
@@ -0,0 +1,36 @@
+"""Coarse-proxy failures: the frame swallowed by vegetation (coarse > 60% vegetation where the fine
+proxy has < 30%), or the protagonist visible in the fine proxy (> 1% of the frame) but nearly absent
+from the coarse one (< 20% of that). A window is bad with >= 3 such samples (every 6th frame)."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from ..clip import n_windows, read_meta, read_proxy, window_frames
+from ..geometry import CLS_HERO, CLS_VEGETATION
+
+SAMPLE_EVERY = 6
+MIN_BAD_SAMPLES = 3
+
+
+def frame_bad(clip: Path, i: int, kind: str = "conditions_coarse") -> bool:
+    _, sf = read_proxy(clip, i)
+    _, sc = read_proxy(clip, i, kind)
+    engulfed = (sc == CLS_VEGETATION).mean() > 0.6 and (sf == CLS_VEGETATION).mean() < 0.3
+    hf, hc = (sf == CLS_HERO).mean(), (sc == CLS_HERO).mean()
+    return bool(engulfed or (hf > 0.01 and hc < 0.2 * hf))
+
+
+def scan(clip: Path, kind: str = "conditions_coarse") -> dict:
+    clip = Path(clip)
+    n = read_meta(clip)["frames"]
+    idx = list(range(0, n, SAMPLE_EVERY))
+    bad = {i: frame_bad(clip, i, kind) for i in idx}
+    windows = []
+    for w in range(n_windows(n)):
+        sl = window_frames(w)
+        k = sum(bad[i] for i in idx if sl.start <= i < sl.stop)
+        windows.append(dict(w=w, bad_samples=int(k), bad=k >= MIN_BAD_SAMPLES))
+    return dict(clip=clip.name, frames=n, bad_frames=[i for i in idx if bad[i]],
+                bad_windows=[x["w"] for x in windows if x["bad"]], windows=windows)
diff --git a/rtcwm/data/qc/lists.py b/rtcwm/data/qc/lists.py
new file mode 100644
index 0000000000000000000000000000000000000000..e3eb5c09d1e56898c86e586168d050ec731dbc5c
--- /dev/null
+++ b/rtcwm/data/qc/lists.py
@@ -0,0 +1,49 @@
+"""The data lists training reads (data_lists/, in git).
+
+    clips.json        {"keep": [...], "dropped": {reason: [...]}}   clips that passed render QC
+    bad_windows.json  {clip: [window, ...]}                          windows flagged by any QC report
+
+Maps dropped as a whole: Town02 (no out-and-back routes), Town04 (navmesh under the terrain:
+protagonist and camera underground, or walking on highways).
+"""
+from __future__ import annotations
+
+import json
+
+from ... import paths
+from ..clip import complete, read_meta
+
+DROPPED_MAPS = {"Town02_Opt": "no out-and-back sidewalk routes", "Town04_Opt": "navmesh under the terrain"}
+CHECKS = ("clipping", "coarse")
+
+
+def build() -> dict:
+    keep, dropped = [], {}
+    for d in sorted(paths.DATASET.glob("clip_*")):
+        if not d.is_dir() or not complete(d):
+            continue
+        m = read_meta(d)["map"]
+        if m in DROPPED_MAPS:
+            dropped.setdefault(DROPPED_MAPS[m], []).append(d.name)
+        else:
+            keep.append(d.name)
+    bad = {}
+    for check in CHECKS:
+        for c in keep:
+            f = paths.QC / check / f"{c}.json"
+            if f.exists():
+                ws = json.loads(f.read_text())["bad_windows"]
+                if ws:
+                    bad[c] = sorted(set(bad.get(c, [])) | set(ws))
+    paths.DATA_LISTS.mkdir(exist_ok=True)
+    (paths.DATA_LISTS / "clips.json").write_text(json.dumps(dict(keep=keep, dropped=dropped), indent=1))
+    (paths.DATA_LISTS / "bad_windows.json").write_text(json.dumps(bad, indent=1, sort_keys=True))
+    return dict(keep=len(keep), dropped={k: len(v) for k, v in dropped.items()},
+                bad_windows=sum(len(v) for v in bad.values()))
+
+
+def load():
+    """(kept clip names, {clip: set(bad windows)})."""
+    clips = json.loads((paths.DATA_LISTS / "clips.json").read_text())["keep"]
+    bad = json.loads((paths.DATA_LISTS / "bad_windows.json").read_text())
+    return clips, {c: set(v) for c, v in bad.items()}
diff --git a/rtcwm/data/render/__init__.py b/rtcwm/data/render/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac68950eff4312522c346df4e4f79075a5a9c66b
--- /dev/null
+++ b/rtcwm/data/render/__init__.py
@@ -0,0 +1,9 @@
+"""Render one CARLA clip in the CWM format.
+
+    python -m rtcwm.data.render --port 2000 --map Town03_Opt --seconds 38.9 --seed 20002 --out OUT
+
+A pedestrian protagonist walks out along the paved middle of a sidewalk, turns round, walks back past
+its start and out again (always walking forwards). open-code-world-model's FollowCamera follows it;
+RGB, depth and instance-segmentation cameras share one pose. See rtcwm.data.clip for the outputs.
+Exit code 3 = rendered but rejected by QC (run_clip.sh then retries with a new seed).
+"""
diff --git a/rtcwm/data/render/__main__.py b/rtcwm/data/render/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e3476db8277a90f143b9bc314fd7e497bd8f025d
--- /dev/null
+++ b/rtcwm/data/render/__main__.py
@@ -0,0 +1,154 @@
+"""python -m rtcwm.data.render --port P --map MAP --seconds S --seed N --out OUT"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import sys
+import time
+from pathlib import Path
+
+import av
+import carla
+import numpy as np
+
+from ... import paths
+from . import camera, scene
+from .config import FPS, HFOV, LEG_PER_SECOND, RGB_H, RGB_W, WALK_SPEED, WARMUP_S, WEATHERS
+from .qc import RenderQC, add_clipping
+from .route import RoutePlanner, Steerer
+from .sensors import SensorRig, decode, leg_clip_fraction
+
+
+def parse():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--host", default="localhost")
+    ap.add_argument("--port", type=int, default=2000)
+    ap.add_argument("--map", default="Town10HD_Opt")
+    ap.add_argument("--seconds", type=float, default=38.916666666666664)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--weather", default=None)
+    ap.add_argument("--vehicles", type=int, default=40)
+    ap.add_argument("--walkers", type=int, default=30)
+    ap.add_argument("--out", required=True)
+    return ap.parse_args()
+
+
+def main():
+    a = parse()
+    a.seed += int(os.environ.get("CARLA_SEED_OFFSET", "0"))     # set by run_clip.sh on a QC retry
+    rng = random.Random(a.seed)
+    out = Path(a.out)
+    (out / "conditions").mkdir(parents=True, exist_ok=True)
+    paths.use_ocwm()
+    from compiler import writer
+    from compiler.render import pool
+
+    client = carla.Client(a.host, a.port)
+    client.set_timeout(300.0)
+    world = client.get_world()
+    if not world.get_map().name.endswith(a.map):
+        world = client.load_world(a.map)
+    settings = world.get_settings()
+    settings.synchronous_mode = True
+    settings.fixed_delta_seconds = 1.0 / FPS
+    world.apply_settings(settings)
+    world.set_pedestrians_seed(a.seed)
+    weather = a.weather or rng.choice(WEATHERS)
+    world.set_weather(getattr(carla.WeatherParameters, weather))
+
+    bp = world.get_blueprint_library()
+    traffic, rig, hero_ctrl = scene.Traffic(), None, None
+    try:
+        traffic = scene.spawn_traffic(world, bp, rng, a.vehicles, a.walkers)
+        leg = LEG_PER_SECOND * a.seconds
+        hero, route = RoutePlanner(world, rng).plan(leg, scene.adult_walkers(bp))
+        traffic.actors.append(hero)
+        print(f"route: {route.kind}, {len(route.targets)} targets, {len(route.hidden_ids)} posts/trunks hidden", flush=True)
+        hero_ctrl = world.spawn_actor(bp.find("controller.ai.walker"), carla.Transform(), hero)
+        traffic.controllers.append(hero_ctrl)
+        world.tick()
+        hero_ctrl.start()
+        hero_ctrl.set_max_speed(WALK_SPEED)
+        steer = Steerer(hero_ctrl, hero, route.targets)
+
+        cam_params = camera.sample_params(rng)
+        follow = camera.FollowRig(world, hero, cam_params)
+        rig = SensorRig(world, bp, follow.pose(1.0 / FPS))
+        for _ in range(int(WARMUP_S * FPS)):                # let the boom and the walkers settle
+            rig.place(follow.pose(1.0 / FPS))
+            rig.grab(world.tick())
+
+        n_frames = int(round(a.seconds * FPS))
+        container = av.open(str(out / "rgb.mp4"), "w")
+        stream = container.add_stream("libx264", rate=FPS)
+        stream.width, stream.height, stream.pix_fmt = RGB_W, RGB_H, "yuv420p"
+        stream.options = {"crf": "16", "preset": "medium"}
+        qc, cams, traj, frames = RenderQC(), [], [], []
+        t0 = time.time()
+        for i in range(n_frames):
+            rig.place(follow.pose(1.0 / FPS))
+            imgs = rig.grab(world.tick())
+            hl = hero.get_location()
+            steer.step(i)
+            f = decode(imgs, hl, rig.rgb)
+            for p in stream.encode(av.VideoFrame.from_ndarray(f.rgb, format="rgb24")):
+                container.mux(p)
+            zp, ip = pool(f.z, f.ids)
+            writer.write_frame(out / "conditions", i, zp, ip)
+            qc.add(hero_frac=float((ip == 8).mean()), near_frac=float(((zp > 0) & (zp < 0.6)).mean()),
+                   leg_clip=leg_clip_fraction(f), luma=float(f.rgb.mean()) if i % FPS == 0 else None)
+            ct = rig.rgb.get_transform()
+            cams.append(dict(loc=[ct.location.x, ct.location.y, ct.location.z],
+                             rot=[ct.rotation.pitch, ct.rotation.yaw, ct.rotation.roll]))
+            traj.append([hl.x, hl.y, hl.z])
+            frames.append(scene.frame_record(hero, traffic.walkers))
+            if i % 240 == 0:
+                print(f"frame {i}/{n_frames}  {(i + 1) / (time.time() - t0):.1f} fps", flush=True)
+        for p in stream.encode():
+            container.mux(p)
+        container.close()
+
+        verdict = qc.summary(traj, leg, WALK_SPEED, a.seconds)
+        centre = (float(np.mean(np.array(traj)[:, 0])), float(np.mean(np.array(traj)[:, 1])))
+        (out / "scene_static.json").write_text(json.dumps(scene.scene_static(world, centre, route.hidden_boxes,
+                                                                             traffic.vehicles)))
+        with open(out / "scene_frames.jsonl", "w") as fh:
+            for sf in frames:
+                fh.write(json.dumps(sf) + "\n")
+        fx = (RGB_W / 2) / math.tan(math.radians(HFOV) / 2)
+        (out / "cameras.json").write_text(json.dumps(dict(
+            convention="UE: x forward, y right, z up, metres; rot = pitch, yaw, roll degrees",
+            intrinsics_rgb=dict(w=RGB_W, h=RGB_H, fx=fx, fy=fx, cx=RGB_W / 2, cy=RGB_H / 2, hfov_deg=HFOV),
+            frames=cams)))
+        verdict = add_clipping(verdict, out)
+        print("qc", verdict, flush=True)
+        hero_frac = np.array(qc.hero_frac)
+        (out / "meta.json").write_text(json.dumps(dict(
+            map=a.map, weather=weather, route="outback", seed=a.seed, frames=n_frames, fps=FPS,
+            route_points=[[p.x, p.y, p.z] for p in route.targets], hero_traj=traj, plan="out-back-out", leg_m=leg,
+            route_kind=route.kind, camera="open-code FollowCamera", cam_params=cam_params,
+            start=[route.start.x, route.start.y, route.start.z],
+            hero_visible_frac_median=float(np.median(hero_frac)), hero_visible_frac_min=float(np.min(hero_frac)),
+            render_fps=n_frames / (time.time() - t0), qc=verdict), indent=1))
+        print("done", out, f"{n_frames / (time.time() - t0):.1f} fps")
+        if not verdict["pass"]:
+            sys.exit(3)
+    finally:
+        if rig is not None:
+            rig.stop()
+        for c in traffic.controllers:
+            try:
+                c.stop()
+            except RuntimeError:
+                pass
+        sensors = rig.sensors if rig is not None else []
+        client.apply_batch([carla.command.DestroyActor(x) for x in sensors + traffic.controllers + traffic.actors])
+        settings.synchronous_mode = False
+        world.apply_settings(settings)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/data/render/camera.py b/rtcwm/data/render/camera.py
new file mode 100644
index 0000000000000000000000000000000000000000..afd21afc27fac51722f96fafb96a5b268752ad40
--- /dev/null
+++ b/rtcwm/data/render/camera.py
@@ -0,0 +1,104 @@
+"""Third-person camera: open-code-world-model's FollowCamera, kept out of solid scenery.
+
+FollowCamera (critically damped boom + stiffer aim, subject held at a fixed screen position) is the
+camera that produces test-time proxies, so training clips use it with parameters drawn from the
+ranges open-code calibrated against the released examples. open-code has no walls; here the boom is
+shortened to stay clear of them:
+
+  * rays from the protagonist's head, hips and knees to the camera and to points 0.3 m around it
+    (one head-to-camera ray passed over a low wall and left the camera inside its cap stone)
+  * thin things -- posts, signs, passers-by -- only count right at the camera (they used to yank it
+    in by 1-3 m within one frame, ~450 such pops in dataset_v1; now it neither pops nor passes
+    through them)
+  * the boom length is rate-limited: in at <= 4 m/s, out at <= 1.5 m/s
+  * a shortened boom re-aims at the point FollowCamera was looking at (keeping the long boom's
+    angles put the protagonist out of frame at wall corners during the turn-round)
+"""
+from __future__ import annotations
+
+import math
+
+import carla
+
+from ... import paths
+from ..geometry import carla_to_renderer, renderer_to_carla
+
+HIP_TO_FEET = 0.93                     # walker actor origin is at the hips
+EYE_ABOVE_HIP = 0.75
+RAY_ORIGINS = (0.75, 0.0, -0.45)       # head, hips, knees (relative to the actor origin)
+RAY_SPREAD = 0.3
+BOOM_IN_SPEED, BOOM_OUT_SPEED = 4.0, 1.5
+MIN_BOOM_M = 0.8
+YAW_TAU = 0.25
+
+
+def sample_params(rng) -> dict:
+    return dict(distance=rng.uniform(1.8, 4.0), height=rng.uniform(1.3, 2.5), aim_height=0.875,
+                aim_screen_x=rng.uniform(0.36, 0.58), aim_screen_y=rng.uniform(0.51, 0.57),
+                focal_px=206.6, boom_tau=0.70, zeta=0.50, aim_tau=0.15)
+
+
+class FollowRig:
+    def __init__(self, world, hero, params: dict):
+        paths.use_ocwm()
+        from compiler.follow import FollowCamera
+        self.world, self.hero, self.params = world, hero, params
+        self.follow = FollowCamera(**params)
+        self.yaw = None                # smoothed walking direction
+        self.k = None                  # current boom fraction (1 = full length)
+        L = carla.CityObjectLabel
+        self.thin = {getattr(L, n) for n in ("Poles", "TrafficSigns", "TrafficLight", "Pedestrians", "Rider", "Other")
+                     if hasattr(L, n)}
+
+    def _heading(self, ht, dt):
+        v = self.hero.get_velocity()
+        head = ht.rotation.yaw if math.hypot(v.x, v.y) < 0.3 else math.degrees(math.atan2(v.y, v.x))
+        if self.yaw is None:
+            self.yaw = head
+        self.yaw += ((head - self.yaw + 180.0) % 360.0 - 180.0) * (1.0 - math.exp(-dt / YAW_TAU))
+        return self.yaw
+
+    def _clear_fraction(self, ht, want, d, dist):
+        side = carla.Location(x=-d.y / dist, y=d.x / dist, z=0.0)
+        offs = [carla.Location()] + [side * s for s in (RAY_SPREAD, -RAY_SPREAD)] + \
+               [carla.Location(z=s) for s in (RAY_SPREAD, -RAY_SPREAD)]
+        k = 1.0
+        for oz in RAY_ORIGINS:
+            o = ht.location + carla.Location(z=oz)
+            for off in offs:
+                tgt = want + off
+                seg = o.distance(tgt)
+                for h in self.world.cast_ray(o, tgt):
+                    t = h.location.distance(o)
+                    if t > 0.4 and (h.label not in self.thin or t > seg - 0.6):
+                        k = min(k, max(0.25, (t - 0.35) / seg))
+        return max(k, min(1.0, MIN_BOOM_M / dist))
+
+    def pose(self, dt: float) -> carla.Transform:
+        ht = self.hero.get_transform()
+        feet = ht.location + carla.Location(z=-HIP_TO_FEET)
+        c = self.follow.update(carla_to_renderer([feet.x, feet.y, feet.z]), self._heading(ht, dt), dt)
+        p0 = carla.Location(*renderer_to_carla(c.position))
+        want = p0
+        eye = ht.location + carla.Location(z=EYE_ABOVE_HIP)
+        d = want - eye
+        dist = math.sqrt(d.x ** 2 + d.y ** 2 + d.z ** 2)
+        if dist > 0.5:
+            k = self._clear_fraction(ht, want, d, dist)
+            if self.k is None:         # first frame: no history to rate-limit against
+                self.k = k
+            if k < self.k:
+                self.k = max(k, self.k - BOOM_IN_SPEED * dt / dist)
+            else:
+                self.k = min(k, self.k + BOOM_OUT_SPEED * dt / dist)
+            if self.k < 1.0:
+                want = eye + carla.Location(x=d.x * self.k, y=d.y * self.k, z=d.z * self.k)
+        pitch, yaw = c.pitch_deg, c.yaw_deg
+        if want.distance(p0) > 0.01:
+            pr, yr = math.radians(pitch), math.radians(yaw)
+            f0 = carla.Location(x=math.cos(pr) * math.cos(yr), y=math.cos(pr) * math.sin(yr), z=math.sin(pr))
+            aim = feet + carla.Location(z=self.params["aim_height"])
+            v = (p0 + f0 * p0.distance(aim)) - want
+            yaw = math.degrees(math.atan2(v.y, v.x))
+            pitch = math.degrees(math.atan2(v.z, math.hypot(v.x, v.y)))
+        return carla.Transform(want, carla.Rotation(pitch=pitch, yaw=yaw))
diff --git a/rtcwm/data/render/config.py b/rtcwm/data/render/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7e405a95b06c9fd96feed8381cbdc1a782004b3
--- /dev/null
+++ b/rtcwm/data/render/config.py
@@ -0,0 +1,30 @@
+"""Renderer constants."""
+from __future__ import annotations
+
+import numpy as np
+
+from ..geometry import HFOV_DEG
+
+FPS = 24
+RGB_W, RGB_H = 1344, 768
+SENSOR_W, SENSOR_H = 672, 384          # depth / segmentation, pooled 2x2 to the 336x192 proxy
+HFOV = HFOV_DEG                        # CWM proxy intrinsics: f = 206.6 px at 336x192
+
+# CARLA 0.9.15 semantic tags -> CWM's released vocabulary {0,1,2,4,5,8,9,10}. Everything else
+# (sidewalk, building, terrain, poles, other pedestrians, ...) stays 0; the protagonist is set to 8
+# from the instance channel.
+TAG2CWM = np.zeros(256, np.uint8)
+TAG2CWM[[1, 24]] = 4                   # road, road line
+TAG2CWM[9] = 5                         # vegetation
+TAG2CWM[11] = 1                        # sky
+TAG2CWM[23] = 2                        # water
+TAG2CWM[[14, 15, 16, 17, 18, 19]] = 10  # car, truck, bus, train, motorcycle, bicycle
+TAG_PEDESTRIAN = 12
+TAGS_SOLID = (3, 4, 5, 9, 20)          # building, wall, fence, vegetation, static -- for the leg check
+
+WEATHERS = ["ClearNoon", "CloudyNoon", "WetNoon", "ClearSunset", "CloudySunset",
+            "WetCloudyNoon", "SoftRainNoon"]   # no night presets: CARLA nights render nearly black
+
+WALK_SPEED = 1.4                       # m/s
+LEG_PER_SECOND = 0.36 * WALK_SPEED     # out-leg length per clip second (out, back, out again)
+WARMUP_S = 1.0
diff --git a/rtcwm/data/render/qc.py b/rtcwm/data/render/qc.py
new file mode 100644
index 0000000000000000000000000000000000000000..de89ca4a1b335f40ab12caed549db84d984a24ec
--- /dev/null
+++ b/rtcwm/data/render/qc.py
@@ -0,0 +1,62 @@
+"""Render-time quality control.
+
+A clip passes when, for no more than 0.5 s at a time, the protagonist is hidden, the camera is
+pressed into geometry, the protagonist's legs are inside solid scenery, it is buried in the ground,
+stands in a hedge / planter, or something stands between it and the camera; when no pole passes
+through its body (>= 3 frames); when the route was completed (got well away, came back); and when the
+image is not too dark. The interpenetration checks are rtcwm.data.qc.clipping, run on the written clip.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from ..geometry import longest_run
+from ..qc import clipping
+from .config import FPS
+
+MAX_BAD_S = 0.5
+MIN_LUMA = 60
+MAX_PIERCED_FRAMES = 3
+
+
+class RenderQC:
+    def __init__(self):
+        self.hero_frac, self.near_frac, self.leg_clip, self.luma = [], [], [], []
+
+    def add(self, hero_frac: float, near_frac: float, leg_clip: float, luma: float | None):
+        self.hero_frac.append(hero_frac)
+        self.near_frac.append(near_frac)
+        self.leg_clip.append(leg_clip)
+        if luma is not None:
+            self.luma.append(luma)
+
+    def summary(self, traj, leg: float, speed: float, seconds: float) -> dict:
+        t = np.asarray(traj)
+        d0 = np.linalg.norm(t[:, :2] - t[0, :2], axis=1)
+        third = len(d0) // 3
+        qc = dict(max_hidden_s=longest_run(np.array(self.hero_frac) < 0.003) / FPS,
+                  max_near_frac=float(np.max(self.near_frac)),
+                  max_near_s=longest_run(np.array(self.near_frac) > 0.30) / FPS,
+                  path_m=float(np.linalg.norm(np.diff(t[:, :2], axis=0), axis=1).sum()),
+                  farthest_first_third_m=float(d0[:third + 1].max()),
+                  closest_return_m=float(d0[third:2 * third + 1].min()),
+                  hero_z_range_m=float(t[:, 2].max() - t[:, 2].min()),
+                  mean_luma=float(np.mean(self.luma)),
+                  max_clip_s=longest_run(np.array(self.leg_clip) > 0.25) / FPS)
+        qc["route_ok"] = bool(qc["path_m"] >= 0.6 * speed * seconds and qc["farthest_first_third_m"] >= 0.5 * leg
+                              and qc["closest_return_m"] <= 6.0)
+        qc["pass"] = bool(qc["route_ok"] and qc["max_hidden_s"] <= MAX_BAD_S and qc["max_near_s"] <= MAX_BAD_S
+                          and qc["hero_z_range_m"] <= 0.3 and qc["mean_luma"] >= MIN_LUMA and qc["max_clip_s"] <= MAX_BAD_S)
+        return qc
+
+
+def add_clipping(qc: dict, clip_dir) -> dict:
+    """Run the interpenetration scan on the written clip and fold it into the verdict."""
+    rep = clipping.scan(clip_dir)
+    qc["max_buried_s"] = rep["max_buried_run"] / FPS
+    qc["max_occluded_s"] = rep["max_occluded_run"] / FPS
+    qc["max_inside_s"] = rep["max_inside_run"] / FPS
+    qc["max_pierced_frames"] = rep["max_pierced_run"]
+    qc["pass"] = bool(qc["pass"] and qc["max_buried_s"] <= MAX_BAD_S and qc["max_occluded_s"] <= MAX_BAD_S
+                      and qc["max_inside_s"] <= MAX_BAD_S and qc["max_pierced_frames"] < MAX_PIERCED_FRAMES)
+    return qc
diff --git a/rtcwm/data/render/route.py b/rtcwm/data/render/route.py
new file mode 100644
index 0000000000000000000000000000000000000000..dcf5b210a6d77092b8ea3b25fd6d19cedfcbe8b8
--- /dev/null
+++ b/rtcwm/data/render/route.py
@@ -0,0 +1,185 @@
+"""The protagonist's out -> back -> out walk.
+
+The walk follows the paved middle of one sidewalk lane: out to a turn point b, back past the start,
+and out again to b2 beyond b (so the last leg never stalls at its end). Targets every 2 m along that
+line are fed to CARLA's walker AI, which routes round street furniture. Posts, signs and tree trunks
+standing on the line are switched off for the clip.
+
+History of the rules (dataset_v1 failures they fix):
+  * walker AI given only far targets took the grass verge, where lamp posts stand, and walked
+    through them (~140 windows)                                   -> dense targets on the paving
+  * CARLA's sidewalk lane often spans the verge (centre line on the grass) and at a block's end
+    next() runs on across the road                                -> paved_line()
+  * Town04 navmesh lies under the terrain (walker rendered underground) -> on_ground(); Town04 is
+    dropped from the dataset anyway
+  * the AI's own arrival radius is larger than ours: on "arriving" it wanders to a random place
+                                                                  -> Steerer re-issues targets
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+import carla
+import numpy as np
+
+TARGET_STEP_M = 2.0
+SWITCH_RADIUS_M = 4.0
+REISSUE_EVERY = 48                     # frames
+HIDE_RADIUS_M = 0.9
+MAX_ATTEMPTS = 300
+NAVMESH_FALLBACK_AFTER = 100           # attempts; maps with almost no sidewalk lanes (Town02/04)
+
+
+@dataclass
+class Route:
+    kind: str                          # "sidewalk" (paved-line route) or "navmesh" (fallback)
+    start: carla.Location
+    b: carla.Location
+    b2: carla.Location
+    leg: float
+    targets: list = field(default_factory=list)
+    line: list = field(default_factory=list)
+    hidden_ids: list = field(default_factory=list)
+    hidden_boxes: list = field(default_factory=list)    # box centres of switched-off objects
+
+
+class RoutePlanner:
+    def __init__(self, world, rng):
+        self.world, self.rng, self.map = world, rng, world.get_map()
+        L = carla.CityObjectLabel
+        self.L = L
+        self.see_through = {getattr(L, n) for n in ("Pedestrians", "Rider", "Vegetation", "Poles", "TrafficSigns",
+                                                    "TrafficLight", "Other") if hasattr(L, n)}
+        self.all_nav = [p for p in (world.get_random_location_from_navigation() for _ in range(4000))
+                        if p is not None and self.on_ground(p)]
+        pool = [p for p in self.all_nav if self.on_sidewalk(p)]
+        # start and turn points must lie on a sidewalk: the navmesh also covers lawns and planters
+        self.pool = pool if len(pool) >= 50 else self.all_nav
+
+    # --- probes
+    def on_sidewalk(self, p) -> bool:
+        w = self.map.get_waypoint(p, project_to_road=True, lane_type=carla.LaneType.Sidewalk)
+        return w is not None and w.transform.location.distance(p) < 1.5
+
+    def on_ground(self, p) -> bool:
+        g = self.world.ground_projection(p + carla.Location(z=3.0), 6.0)
+        return g is not None and abs(g.location.z - p.z) < 0.3
+
+    def ground_label(self, q):
+        top = q + carla.Location(z=1.5)
+        hits = [h for h in self.world.cast_ray(top, q + carla.Location(z=-1.5)) if h.label not in self.see_through]
+        return min(hits, key=lambda h: h.location.distance(top)).label if hits else None
+
+    def snap(self, q, tol: float = 1.0):
+        c = min(self.all_nav, key=lambda p: (p.x - q.x) ** 2 + (p.y - q.y) ** 2)
+        return c if math.hypot(c.x - q.x, c.y - q.y) < tol else None
+
+    def paved_line(self, w, length: float, step: float = TARGET_STEP_M):
+        """Points every `step` m along a sidewalk lane, each moved across the lane to the middle of
+        its paved part. None if a point has no paving under it or the lane jumps (a junction)."""
+        pts, run_m, prev = [], 0.0, None
+        while True:
+            loc, right = w.transform.location, w.transform.get_right_vector()
+            if prev is not None and loc.distance(prev) > 1.5 * step:
+                return None
+            half = 0.5 * w.lane_width
+            offs = np.arange(-half + 0.2, half - 0.2 + 1e-6, 0.3)
+            paved = [self.ground_label(loc + right * float(o)) == self.L.Sidewalks for o in offs]
+            if not any(paved):
+                return None
+            best, cur, bs, cs = 0, 0, 0, 0                      # longest run of paved samples
+            for i, p in enumerate(paved):
+                cur, cs = (cur + 1, cs) if p else (0, i + 1)
+                if cur > best:
+                    best, bs = cur, cs
+            pts.append(loc + right * float(np.mean(offs[bs:bs + best])))
+            prev = loc
+            if run_m >= length:
+                return pts
+            nx = w.next(step)
+            if not nx:
+                return None
+            w, run_m = nx[0], run_m + step
+
+    # --- planning
+    def _spawn(self, bps, start, b):
+        hb = self.rng.choice(bps)
+        if hb.has_attribute("is_invincible"):
+            hb.set_attribute("is_invincible", "true")
+        yaw = math.degrees(math.atan2(b.y - start.y, b.x - start.x))
+        return self.world.try_spawn_actor(hb, carla.Transform(start + carla.Location(z=1.0), carla.Rotation(yaw=yaw)))
+
+    def plan(self, leg: float, walker_bps):
+        """Spawn the protagonist on a valid route. Returns (hero, Route)."""
+        for attempt in range(MAX_ATTEMPTS):
+            start = self.rng.choice(self.pool)
+            w0 = self.map.get_waypoint(start, project_to_road=True, lane_type=carla.LaneType.Sidewalk)
+            nb = w0.next(leg) if w0 else []
+            nb2 = w0.next(1.4 * leg) if w0 else []
+            if (not nb or not nb2) and attempt >= NAVMESH_FALLBACK_AFTER:
+                far = [q for q in self.all_nav if 0.75 * leg <= math.hypot(q.x - start.x, q.y - start.y) <= leg]
+                if not far:
+                    continue
+                b = self.rng.choice(far)
+                ext = carla.Location(x=start.x + 1.4 * (b.x - start.x), y=start.y + 1.4 * (b.y - start.y), z=b.z)
+                b2 = min(self.all_nav, key=lambda q: (q.x - ext.x) ** 2 + (q.y - ext.y) ** 2)
+                hero = self._spawn(walker_bps, start, b)
+                if hero:
+                    return hero, Route("navmesh", start, b, b2, leg, targets=[b, start, b2])
+                continue
+            if not nb or not nb2:
+                continue
+            line = self.paved_line(w0, 1.4 * leg)       # the whole walk stays on this lane's paving
+            if line is None:
+                continue
+            n_out = max(1, int(leg / TARGET_STEP_M))
+            start, b, b2 = self.snap(line[0]), self.snap(line[n_out]), self.snap(line[-1])
+            if start is None or b is None or b2 is None:
+                continue
+            if math.hypot(b.x - start.x, b.y - start.y) < 0.6 * leg:     # lane bends back on itself
+                continue
+            hero = self._spawn(walker_bps, start, b)
+            if hero:
+                route = Route("sidewalk", start, b, b2, leg, targets=[b, start, b2], line=line)
+                out_pts = [p for p in (self.snap(q) for q in line[1:n_out + 1]) if p is not None]
+                out2_pts = [p for p in (self.snap(q) for q in line[1:]) if p is not None]
+                if len(out_pts) >= 3 and len(out2_pts) > len(out_pts):
+                    route.targets = out_pts + list(reversed(out_pts[:-1])) + [start] + out2_pts
+                self.hide_obstacles(route)
+                return hero, route
+        raise RuntimeError("could not spawn a protagonist with an out-and-back route")
+
+    def hide_obstacles(self, route: Route):
+        """Switch off posts, signs, traffic lights and tree trunks standing on the walking line."""
+        xy = np.array([[p.x, p.y] for p in route.line])
+        for lab in ("Poles", "TrafficSigns", "TrafficLight", "Vegetation"):
+            if not hasattr(self.L, lab):
+                continue
+            for ob in self.world.get_environment_objects(getattr(self.L, lab)):
+                bb = ob.bounding_box
+                if lab == "Vegetation" and bb.extent.z < 1.2:
+                    continue                                    # hedges / bushes: not on the paving
+                if float(np.min(np.hypot(xy[:, 0] - bb.location.x, xy[:, 1] - bb.location.y))) < HIDE_RADIUS_M:
+                    route.hidden_ids.append(ob.id)
+                    route.hidden_boxes.append([bb.location.x, bb.location.y, bb.location.z])
+        if route.hidden_ids:
+            self.world.enable_environment_objects(route.hidden_ids, False)
+
+
+class Steerer:
+    """Feeds the route targets to the walker AI: the next target from SWITCH_RADIUS_M out, and the
+    current one re-issued every REISSUE_EVERY frames."""
+
+    def __init__(self, controller, hero, targets):
+        self.ctrl, self.hero, self.targets, self.i = controller, hero, targets, 0
+        self.ctrl.go_to_location(targets[0])
+
+    def step(self, frame: int):
+        p = self.hero.get_location()
+        t = self.targets[self.i]
+        if self.i < len(self.targets) - 1 and math.hypot(p.x - t.x, p.y - t.y) < SWITCH_RADIUS_M:
+            self.i += 1
+            self.ctrl.go_to_location(self.targets[self.i])
+        elif frame % REISSUE_EVERY == 0:
+            self.ctrl.go_to_location(t)
diff --git a/rtcwm/data/render/scene.py b/rtcwm/data/render/scene.py
new file mode 100644
index 0000000000000000000000000000000000000000..076a95a22f99d4e098c77aa981ad72d10f47bdfa
--- /dev/null
+++ b/rtcwm/data/render/scene.py
@@ -0,0 +1,100 @@
+"""Background traffic and the scene description saved for the coarse proxy."""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+import carla
+
+# level bounding boxes exported for the coarse proxy: CARLA label -> CWM class
+STATIC_LABELS = {"Buildings": 0, "Walls": 0, "Fences": 0, "Poles": 0, "TrafficLight": 0, "TrafficSigns": 0,
+                 "Static": 0, "Dynamic": 0, "Bridge": 0, "GuardRail": 0, "Vegetation": 5, "Water": 2}
+SCENE_RADIUS_M = 300.0
+
+
+@dataclass
+class Traffic:
+    actors: list = field(default_factory=list)
+    controllers: list = field(default_factory=list)
+    vehicles: list = field(default_factory=list)
+    walkers: list = field(default_factory=list)
+
+
+def spawn_traffic(world, bp, rng, n_vehicles: int, n_walkers: int) -> Traffic:
+    """Parked vehicles (physics off: no traffic manager, it segfaults libcarla on some nodes, and a
+    static street is what a revisit dataset wants) and AI-driven background pedestrians."""
+    t = Traffic()
+    spawns = world.get_map().get_spawn_points()
+    rng.shuffle(spawns)
+    vbps = [b for b in bp.filter("vehicle.*") if int(b.get_attribute("number_of_wheels")) == 4]
+    for sp in spawns[:n_vehicles]:
+        v = world.try_spawn_actor(rng.choice(vbps), sp)
+        if v:
+            v.set_simulate_physics(False)
+            t.actors.append(v)
+            t.vehicles.append(v)
+    wbps = bp.filter("walker.pedestrian.*")
+    ctrl_bp = bp.find("controller.ai.walker")
+    for _ in range(n_walkers):
+        loc = world.get_random_location_from_navigation()
+        if loc is None:
+            continue
+        w = world.try_spawn_actor(rng.choice(wbps), carla.Transform(loc + carla.Location(z=1.0)))
+        if w:
+            t.actors.append(w)
+            t.walkers.append(w)
+            t.controllers.append(world.spawn_actor(ctrl_bp, carla.Transform(), w))
+    world.tick()
+    for c in t.controllers:
+        c.start()
+        c.go_to_location(world.get_random_location_from_navigation())
+        c.set_max_speed(1.0 + rng.random() * 0.6)
+    return t
+
+
+def adult_walkers(bp):
+    return [b for b in bp.filter("walker.pedestrian.*")
+            if not (b.has_attribute("age") and b.get_attribute("age").as_str() == "child")]
+
+
+def frame_record(hero, walkers) -> dict:
+    """Per-frame scene state for the coarse proxy: protagonist bones, background walker poses."""
+    bones = {bt.name: [round(bt.world.location.x, 3), round(bt.world.location.y, 3), round(bt.world.location.z, 3)]
+             for bt in hero.get_bones().bone_transforms}
+    others = []
+    for w in walkers:
+        wt = w.get_transform()
+        others.append([round(wt.location.x, 2), round(wt.location.y, 2), round(wt.location.z, 2), round(wt.rotation.yaw, 1)])
+    return dict(bones=bones, walkers=others)
+
+
+def scene_static(world, centre, hidden, vehicles) -> dict:
+    """Level boxes, driving-lane waypoints and parked vehicles within SCENE_RADIUS_M of `centre`.
+    Objects switched off for the walk (`hidden`: their box centres) are left out."""
+    L = carla.CityObjectLabel
+    cx, cy = centre
+    boxes = []
+    for name, cls in STATIC_LABELS.items():
+        if not hasattr(L, name):
+            continue
+        for bb in world.get_level_bbs(getattr(L, name)):
+            if math.hypot(bb.location.x - cx, bb.location.y - cy) > SCENE_RADIUS_M:
+                continue
+            if any(abs(bb.location.x - h[0]) < 0.05 and abs(bb.location.y - h[1]) < 0.05
+                   and abs(bb.location.z - h[2]) < 0.05 for h in hidden):
+                continue
+            boxes.append(dict(cls=cls, label=name, c=[bb.location.x, bb.location.y, bb.location.z],
+                              e=[bb.extent.x, bb.extent.y, bb.extent.z],
+                              r=[bb.rotation.pitch, bb.rotation.yaw, bb.rotation.roll]))
+    lanes = []
+    for w in world.get_map().generate_waypoints(2.0):
+        p = w.transform.location
+        if math.hypot(p.x - cx, p.y - cy) <= SCENE_RADIUS_M:
+            lanes.append([p.x, p.y, p.z, w.transform.rotation.yaw, w.lane_width])
+    vboxes = []
+    for v in vehicles:
+        bb, vt = v.bounding_box, v.get_transform()
+        c = vt.transform(bb.location)
+        vboxes.append(dict(c=[c.x, c.y, c.z], e=[bb.extent.x, bb.extent.y, bb.extent.z],
+                           r=[vt.rotation.pitch, vt.rotation.yaw, vt.rotation.roll]))
+    return dict(boxes=boxes, lanes=lanes, vehicles=vboxes)
diff --git a/rtcwm/data/render/sensors.py b/rtcwm/data/render/sensors.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f39d4d52fa47a734149bba28af9bf83fc017f15
--- /dev/null
+++ b/rtcwm/data/render/sensors.py
@@ -0,0 +1,113 @@
+"""RGB / depth / instance-segmentation cameras on one pose, and decoding their frames."""
+from __future__ import annotations
+
+import math
+import queue
+from dataclasses import dataclass
+
+import numpy as np
+
+from .config import HFOV, RGB_H, RGB_W, SENSOR_H, SENSOR_W, TAG2CWM, TAG_PEDESTRIAN, TAGS_SOLID
+
+RGB, DEPTH, SEG = "sensor.camera.rgb", "sensor.camera.depth", "sensor.camera.instance_segmentation"
+
+
+class SensorRig:
+    def __init__(self, world, bp, initial_tf):
+        self.sensors, self.queues = [], {}
+        self.rgb = self._add(world, bp, initial_tf, RGB, RGB_W, RGB_H, motion_blur_intensity=0.0)
+        self._add(world, bp, initial_tf, DEPTH, SENSOR_W, SENSOR_H)
+        self._add(world, bp, initial_tf, SEG, SENSOR_W, SENSOR_H)
+
+    def _add(self, world, bp, tf, kind, w, h, **attrs):
+        b = bp.find(kind)
+        b.set_attribute("image_size_x", str(w))
+        b.set_attribute("image_size_y", str(h))
+        b.set_attribute("fov", f"{HFOV:.4f}")
+        for k, v in attrs.items():
+            if b.has_attribute(k):
+                b.set_attribute(k, str(v))
+        s = world.spawn_actor(b, tf)
+        q = queue.Queue()
+        s.listen(q.put)
+        self.sensors.append(s)
+        self.queues[kind] = q
+        return s
+
+    def place(self, tf):
+        for s in self.sensors:
+            s.set_transform(tf)
+
+    def grab(self, frame: int, timeout: float = 120.0) -> dict:
+        got = {}
+        for kind, q in self.queues.items():
+            while True:
+                img = q.get(timeout=timeout)
+                if img.frame == frame:
+                    got[kind] = img
+                    break
+        return got
+
+    def stop(self):
+        for s in self.sensors:
+            s.stop()
+
+
+@dataclass
+class Frame:
+    rgb: np.ndarray            # [H,W,3] uint8
+    z: np.ndarray              # [SENSOR_H,SENSOR_W] planar depth (m), sky -> 0, clipped at 0.3
+    ids: np.ndarray            # CWM semantic ids at sensor resolution
+    tag: np.ndarray            # raw CARLA tags
+    hero: np.ndarray           # protagonist mask
+
+
+def decode(imgs: dict, hero_loc, rgb_sensor) -> Frame:
+    bgra = np.frombuffer(imgs[RGB].raw_data, np.uint8).reshape(RGB_H, RGB_W, 4)
+    d = np.frombuffer(imgs[DEPTH].raw_data, np.uint8).reshape(SENSOR_H, SENSOR_W, 4).astype(np.float64)
+    z = (d[..., 2] + d[..., 1] * 256.0 + d[..., 0] * 65536.0) / (256.0 ** 3 - 1) * 1000.0   # 1000 m far
+    s = np.frombuffer(imgs[SEG].raw_data, np.uint8).reshape(SENSOR_H, SENSOR_W, 4)
+    tag = s[..., 2]
+    inst = s[..., 1].astype(np.int32) + 256 * s[..., 0].astype(np.int32)
+    ids = TAG2CWM[tag].copy()
+    hero = hero_mask(tag, inst, hero_loc, np.array(rgb_sensor.get_transform().get_inverse_matrix()))
+    ids[hero] = 8
+    sky = (ids == 1) | (z >= 999.0)
+    ids[sky] = 1
+    z = np.where(sky, 0.0, np.clip(z, 0.3, None)).astype(np.float32)
+    return Frame(np.ascontiguousarray(bgra[..., 2::-1]), z, ids, tag, hero)
+
+
+def hero_mask(tag, inst, hero_loc, w2c) -> np.ndarray:
+    """The instance channel is an internal object id, not the actor id: project the protagonist's
+    body into the image and take the pedestrian instance found there (majority vote)."""
+    fpx = (SENSOR_W / 2) / math.tan(math.radians(HFOV) / 2)
+    votes = {}
+    for dz in (-0.4, 0.0, 0.4):
+        pc = w2c @ np.array([hero_loc.x, hero_loc.y, hero_loc.z + dz, 1.0])
+        if pc[0] <= 0.3:
+            continue
+        u = int(SENSOR_W / 2 + fpx * pc[1] / pc[0])
+        v = int(SENSOR_H / 2 - fpx * pc[2] / pc[0])
+        y0, y1, x0, x1 = max(v - 6, 0), min(v + 7, SENSOR_H), max(u - 6, 0), min(u + 7, SENSOR_W)
+        if y0 >= y1 or x0 >= x1:
+            continue
+        win_t, win_i = tag[y0:y1, x0:x1], inst[y0:y1, x0:x1]
+        for key in win_i[win_t == TAG_PEDESTRIAN].ravel():
+            votes[int(key)] = votes.get(int(key), 0) + 1
+    if not votes:
+        return np.zeros_like(tag, bool)
+    return (tag == TAG_PEDESTRIAN) & (inst == max(votes, key=votes.get))
+
+
+def leg_clip_fraction(f: Frame) -> float:
+    """Fraction of the protagonist's leg box covered by solid scenery (building, wall, fence,
+    vegetation, props) at its own depth. Ground classes are excluded -- they sit at foot depth."""
+    ys, xs = np.nonzero(f.hero)
+    if len(ys) <= 80:
+        return 0.0
+    y0, y1, x0, x1 = ys.min(), ys.max() + 1, xs.min(), xs.max() + 1
+    legs = slice(y0 + int(0.55 * (y1 - y0)), y1)
+    hz = float(np.median(f.z[ys, xs]))
+    solid = np.isin(f.tag[legs, x0:x1], TAGS_SOLID) & (np.abs(f.z[legs, x0:x1] - hz) < 0.7)
+    return float(solid.mean())
diff --git a/rtcwm/eval/__init__.py b/rtcwm/eval/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..377a31bb74c8ea0ace4f6d45e6756f5b1fa8e46e
--- /dev/null
+++ b/rtcwm/eval/__init__.py
@@ -0,0 +1 @@
+"""Evaluation: chunk-autoregressive generation of held-out windows against the CARLA ground truth."""
diff --git a/rtcwm/eval/compare.py b/rtcwm/eval/compare.py
new file mode 100644
index 0000000000000000000000000000000000000000..33c88dff7141dd35d1fd76a4f768bed26d3035db
--- /dev/null
+++ b/rtcwm/eval/compare.py
@@ -0,0 +1,56 @@
+"""Side-by-side comparison of evaluation runs: [run A gen | run B gen | ... | GT] per sample, plus a
+summary table of pixel MSE.
+
+    python -m rtcwm.eval.compare --runs EVAL_DIR_A EVAL_DIR_B --labels "step 0" "step 49" --out DIR
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from .video import read_mp4, write_mp4
+
+W = 672
+
+
+def _label(im, text):
+    im = im.copy()
+    cv2.putText(im, text, (10, 30), 0, 0.9, (0, 0, 0), 5)
+    cv2.putText(im, text, (10, 30), 0, 0.9, (255, 255, 255), 2)
+    return im
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--runs", nargs="+", type=Path, required=True)
+    ap.add_argument("--labels", nargs="+", required=True)
+    ap.add_argument("--out", type=Path, required=True)
+    a = ap.parse_args()
+    assert len(a.runs) == len(a.labels)
+    a.out.mkdir(parents=True, exist_ok=True)
+    names = sorted(p.name.replace("_gen_vs_gt.mp4", "") for p in a.runs[0].glob("*_gen_vs_gt.mp4"))
+    table = []
+    for n in names:
+        vids = [read_mp4(r / f"{n}_gen_vs_gt.mp4") for r in a.runs]
+        half = vids[0].shape[2] // 2
+        cols = [v[..., :half, :] for v in vids] + [vids[0][..., half:, :]]
+        labels = list(a.labels) + ["ground truth"]
+        h = int(cols[0].shape[1] * W / half)
+        frames = [np.concatenate([_label(cv2.resize(c[i], (W, h)), lab) for c, lab in zip(cols, labels)], 1)
+                  for i in range(min(len(c) for c in cols))]
+        write_mp4(a.out / f"{n}_compare.mp4", np.stack(frames), crf=22)
+        row = dict(sample=n)
+        for r, lab in zip(a.runs, a.labels):
+            f = r / f"{n}.json"
+            row[lab] = json.loads(f.read_text())["pixel_mse"] if f.exists() else None
+        table.append(row)
+        print(json.dumps(row), flush=True)
+    (a.out / "summary.json").write_text(json.dumps(table, indent=1))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/eval/evaluate.py b/rtcwm/eval/evaluate.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3bcf0a8af85430ece4eab56357bf0d5cae110b0
--- /dev/null
+++ b/rtcwm/eval/evaluate.py
@@ -0,0 +1,106 @@
+"""Stream held-out clips chunk by chunk and compare with the CARLA ground truth.
+
+    python -m rtcwm.eval.evaluate --lora LORA --out DIR --clips clip_0244 [...] [--chunks N]
+           [--steps 35] [--rule euler|x0] [--shift S]          (distilled students: --steps 4 --rule x0 --shift 5)
+
+Each clip is generated from nothing (no keyframe) as one stream of --chunks 17-frame chunks (default: all
+54, 918 frames / 38 s), every chunk in its window of the 5 chunks before it (rtcwm.model.sampler).
+Writes DIR/<clip>_stream_gen_vs_gt.mp4 (generated | GT, full resolution), DIR/<clip>_latents.pt and DIR/<clip>_stream.json:
+pixel MSE overall and per chunk, the luma jump at chunk boundaries vs inside chunks, seconds per chunk.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from cwm_h3_inference.constants import VIDEO_SHIFT
+
+from .. import paths
+from ..model.layout import CLIP_FRAMES
+from ..model.loading import load_for_inference
+from ..model.sampler import generate_stream
+from ..train.dataset import load_window, silence, stream_chunks
+from .video import decode, load_vae, read_mp4, write_mp4
+
+
+def luma(frames: np.ndarray) -> np.ndarray:
+    f = frames.astype(np.float32)
+    return (0.299 * f[..., 0] + 0.587 * f[..., 1] + 0.114 * f[..., 2]).mean(axis=(1, 2))
+
+
+def report(gen: np.ndarray, gt: np.ndarray) -> dict:
+    err = ((gen.astype(np.float32) - gt.astype(np.float32)) ** 2).mean(axis=(1, 2, 3))
+    jumps, gt_jumps = np.abs(np.diff(luma(gen))), np.abs(np.diff(luma(gt)))
+    boundary = np.zeros(len(jumps), bool)
+    boundary[CLIP_FRAMES - 1::CLIP_FRAMES] = True                  # frame 17k-1 -> 17k
+    n = len(gen) // CLIP_FRAMES
+    return dict(frames=int(len(gen)), pixel_mse=float(err.mean()),
+                chunk_mse=[float(err[CLIP_FRAMES * k:CLIP_FRAMES * (k + 1)].mean()) for k in range(n)],
+                boundary_luma_jump=float(jumps[boundary].mean()), inner_luma_jump=float(jumps[~boundary].mean()),
+                gt_boundary_luma_jump=float(gt_jumps[boundary].mean()), max_luma_jump=float(jumps.max()))
+
+
+def check_span(lora):
+    """The LoRA's training span (its metadata; 6 if it predates the setting) must be the sampler's."""
+    from safetensors import safe_open
+    from ..model.layout import SPAN_CHUNKS
+    with safe_open(str(lora), "pt") as f:
+        trained = int((f.metadata() or {}).get("span_chunks", "6"))
+    assert trained == SPAN_CHUNKS, f"{lora} was trained with span {trained}; set RTCWM_SPAN_CHUNKS={trained}"
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--lora", type=Path, default=paths.CWM_LORA)
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--clips", nargs="+", required=True)
+    ap.add_argument("--chunks", type=int, default=0, help="0 = the whole stream")
+    ap.add_argument("--steps", type=int, default=35)
+    ap.add_argument("--rule", choices=["euler", "x0"], default="euler", help="x0: few-step (distilled) sampling")
+    ap.add_argument("--shift", type=float, default=VIDEO_SHIFT, help="flow shift of the sampling schedule")
+    ap.add_argument("--history-noise", type=float, default=0.0, help="noise level of the history chunks")
+    ap.add_argument("--picture1", choices=["native", "2048"], default="native", help="Picture 1 resolution the LoRA was trained with")
+    ap.add_argument("--kv-reuse", action="store_true", help="carry the history's keys/values across chunks")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--fp8", action="store_true", help="the blocks' large linears in FP8 (Hopper)")
+    ap.add_argument("--history-from-last-step", action=argparse.BooleanOptionalAction, default=True,
+                    help="kv-reuse: cache each chunk from its last denoising step (one forward fewer per chunk; "
+                         "--no-history-from-last-step: one more clean forward)")
+    a = ap.parse_args()
+    dev = torch.device("cuda", 0)
+    a.out.mkdir(parents=True, exist_ok=True)
+    check_span(a.lora)
+    from ..model import sampler
+    sampler.HISTORY_FROM_LAST_STEP = a.history_from_last_step
+    model = load_for_inference(a.lora, dev, fp8=a.fp8)
+    audio, results = silence(), {}
+    for clip in a.clips:
+        n = a.chunks or stream_chunks(clip)
+        secs = []
+        t0 = time.time()
+        lat = generate_stream(model, lambda j: load_window(clip, j, audio, start=0, with_target=False, res=a.picture1), n, a.steps,
+                              torch.Generator().manual_seed(a.seed), dev, a.rule, a.shift, a.history_noise, a.kv_reuse,
+                              on_chunk=lambda k, s: (secs.append(round(s, 1)),
+                                                     print(f"{clip} chunk {k + 1}/{n} {s:.1f}s", flush=True)))
+        results[clip] = (lat.cpu(), secs, time.time() - t0)
+    del model
+    torch.cuda.empty_cache()
+    vae = load_vae(dev)
+    for clip, (lat, secs, total) in results.items():
+        gen = decode(vae, lat, dev)
+        gt = read_mp4(paths.DATASET / clip / "rgb.mp4")[:len(gen)]
+        rep = dict(clip=clip, lora=str(a.lora), steps=a.steps, rule=a.rule, history_noise=a.history_noise, history_from_last_step=a.history_from_last_step, kv_reuse=a.kv_reuse, chunks=len(secs), **report(gen, gt),
+                   seconds_per_chunk=secs, gen_seconds=round(total, 1))
+        (a.out / f"{clip}_stream.json").write_text(json.dumps(rep, indent=1))
+        torch.save(lat, a.out / f"{clip}_latents.pt")
+        write_mp4(a.out / f"{clip}_stream_gen_vs_gt.mp4", np.concatenate([gen, gt], 2), crf=20)
+        print(json.dumps({k: v for k, v in rep.items() if k not in ("chunk_mse", "seconds_per_chunk")}), flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/eval/open_world.py b/rtcwm/eval/open_world.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c5754c31d56772aaad60355d2743c43caf52395
--- /dev/null
+++ b/rtcwm/eval/open_world.py
@@ -0,0 +1,185 @@
+"""Stream open-world examples -- the released ones (code-world-model release/examples/example_XX: Picture 1, a
+Depth + Semantic-ID proxy, per-window captions) or open-code-world-model takes (examples/gallery.py: a released
+first frame, a proxy compiled by the same renderer and rig as our CARLA proxies). There is no ground truth: a
+qualitative test outside CARLA.
+
+    python -m rtcwm.eval.open_world prepare --examples DIR/example_01 [...] [--chunks 24]    # one GPU
+    python -m rtcwm.eval.open_world generate --lora LORA --out DIR --examples example_01 [...] [--chunks 24]
+    python -m rtcwm.eval.open_world release --out DIR --examples example_01 [...] [--chunks 24]
+
+prepare writes each example in the training-cache layout (train_cache/ow_example_XX/), so
+rtcwm.train.dataset.load_window streams it like a CARLA clip:
+    stream.safetensors {"duv"}   proxy latents of the first --chunks 17-frame clips
+    anchorn_c00 / staticn_c00    Picture 1 at 1344x768 (the release's first target frame) and the Qwen states
+                                 of the static caption
+The static caption is the example's first-window "Appearance and scene" section -- the starting state --
+without its action timeline (real time: the text tells nothing of the future).
+generate writes <example>_latents.pt and <example>_stream.mp4: generated | proxy (Semantic-ID colours), full
+resolution.
+release writes DIR/<example>/config.json: the example's own release configuration cut to the windows that
+cover --chunks chunks, for the released (bidirectional, windowed) CWM sampler -- python -m
+cwm_h3_inference prepare / generate --config ...; it keeps the full per-window captions, timeline included.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from cwm_h3_inference.constants import TARGET_HEIGHT, TARGET_WIDTH
+
+from .. import paths
+from ..data.anchors import cache_paths
+from ..model.layout import CLIP_FRAMES
+
+TIMELINE = "Protagonist action timeline:"
+PALETTE = np.array([[0, 0, 0], [135, 206, 235], [30, 90, 200], [150, 120, 80], [90, 90, 90], [60, 160, 60],
+                    [200, 120, 80], [220, 200, 60], [230, 40, 40], [200, 80, 200], [40, 220, 220], [250, 250, 250]],
+                   dtype=np.uint8)                                 # Semantic-ID classes 0..11 (INFERENCE.md)
+
+
+def cache_name(example: str) -> str:
+    return "ow_" + Path(example).name
+
+
+def static_caption(prompt: str) -> str:
+    """A window prompt cut to its "Appearance and scene:" section -- the starting state, no timeline (and no
+    instruction preamble, which open-code-world-model's captions carry)."""
+    head = "Appearance and scene:"
+    return (head + prompt.split(head, 1)[1] if head in prompt else prompt).split(TIMELINE)[0].strip()
+
+
+def example_inputs(ex: Path):
+    """(Picture 1, condition directory, first-window prompt) of a released example (config.json) or of an
+    open-code-world-model take (first_frame.png, conditions/, prompts/000.txt)."""
+    if (ex / "config.json").exists():
+        cfg = json.loads((ex / "config.json").read_text())
+        return ex / cfg["anchor_image"], ex / cfg["condition_root"], cfg["windows"][0]["prompt"]
+    return ex / "first_frame.png", ex / "conditions", (ex / "prompts" / "000.txt").read_text()
+
+
+def prepare(examples: list[Path], chunks: int):
+    from PIL import Image
+    from musubi_tuner.minimax_h3.text_encoder import load_h3_processor, load_h3_text_encoder
+    from musubi_tuner.minimax_h3.video_vae import encode_video_condition, load_video_vae
+    from musubi_tuner.minimax_h3_generate_video import _prepare_pixels
+    from cwm_h3_inference.cache import _load_anchor
+    from cwm_h3_inference.config import _canonical_caption
+    from cwm_h3_inference.duv import load_duv_frame
+    from ..train.preencode import atomic_save, static_text_states
+    dev = torch.device("cuda:0")
+    vae = load_video_vae(paths.VIDEO_VAE, device=dev, dtype=torch.float16, disable_mmap=False).eval().requires_grad_(False)
+    proc = load_h3_processor(str(paths.PROCESSOR))
+    qwen = load_h3_text_encoder(str(paths.TEXT_ENCODER), processor_path=str(paths.PROCESSOR), device=dev,
+                                dtype=torch.bfloat16).eval()
+    for ex in examples:
+        anchor, cond, prompt = example_inputs(ex)
+        name = cache_name(ex)
+        (paths.CACHE / name).mkdir(parents=True, exist_ok=True)
+        k = min(chunks, len(list(cond.glob("*.semantic_id.png"))) // CLIP_FRAMES)
+        duv = np.stack([load_duv_frame(cond, i, for_qwen=False) for i in range(k * CLIP_FRAMES)])
+        duv = np.concatenate([duv, np.repeat(duv[-1:], CLIP_FRAMES, 0)], 0)     # the encoder drops 3 latents
+        with torch.inference_mode():
+            ref = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(duv)).to(dev, torch.float16))[0]
+        atomic_save(paths.CACHE / name / "stream.safetensors", {"duv": ref.float().cpu()[:, :5 * k].contiguous()},
+                    {"clip_frames": str(CLIP_FRAMES), "vae_clips": str(k), "source": str(ex)})
+        cp = cache_paths(name, 0, "native")
+        picture1 = _load_anchor(anchor, size=(TARGET_WIDTH, TARGET_HEIGHT))
+        Image.fromarray(picture1).save(cp["png"])
+        with torch.inference_mode():
+            lat = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(picture1).unsqueeze(0)).to(dev, torch.float16))
+        atomic_save(cp["latent"], {"latents_ref_000_image": lat[0].float().cpu()})
+        caption = _canonical_caption(static_caption(prompt), f"{name}/static_c00")
+        hidden, tags = static_text_states(proc, qwen, picture1, cp["png"], cond, 0, caption)
+        atomic_save(cp["text"], {"hidden": hidden.to("cpu", torch.bfloat16), "tags": tags.to("cpu", torch.int64)},
+                    {"system_prompt": "w0", "duv": "start frame repeated", "caption": "first window, no timeline"})
+        print(f"{name}: {k} chunks, caption {len(caption.split())} words, Qwen rows {hidden.shape[0]}", flush=True)
+
+
+def proxy_frames(example: Path, n: int) -> np.ndarray:
+    from PIL import Image
+    cond = example_inputs(example)[1]
+    sem = [np.asarray(Image.open(cond / f"{i:06d}.semantic_id.png")) for i in range(n)]
+    out = PALETTE[np.clip(np.stack(sem), 0, 11)]
+    return np.stack([np.asarray(Image.fromarray(f).resize((TARGET_WIDTH, TARGET_HEIGHT), Image.NEAREST)) for f in out])
+
+
+def generate(a):
+    from cwm_h3_inference.constants import VIDEO_SHIFT
+    from ..model.loading import load_for_inference
+    from ..model.sampler import generate_stream
+    from ..train.dataset import load_window, silence, stream_chunks
+    from .video import decode, load_vae, write_mp4
+    dev = torch.device("cuda", 0)
+    a.out.mkdir(parents=True, exist_ok=True)
+    from .evaluate import check_span
+    check_span(a.lora)
+    from ..model import sampler
+    sampler.HISTORY_FROM_LAST_STEP = a.history_from_last_step
+    model = load_for_inference(a.lora, dev, fp8=a.fp8)
+    audio, results = silence(), {}
+    for ex in a.examples:
+        name = cache_name(ex)
+        n = min(a.chunks, stream_chunks(name))
+        lat = generate_stream(model, lambda j: load_window(name, j, audio, start=0, with_target=False), n, a.steps,
+                              torch.Generator().manual_seed(a.seed), dev, a.rule, a.shift if a.shift else VIDEO_SHIFT,
+                              history_noise=a.history_noise, kv_reuse=a.kv_reuse, on_chunk=lambda k, s: print(f"{name} chunk {k + 1}/{n} {s:.1f}s", flush=True))
+        results[ex] = lat.cpu()
+    del model
+    torch.cuda.empty_cache()
+    vae = load_vae(dev)
+    for ex, lat in results.items():
+        torch.save(lat, a.out / f"{ex}_latents.pt")
+        gen = decode(vae, lat, dev)
+        write_mp4(a.out / f"{ex}_stream.mp4", np.concatenate([gen, proxy_frames(Path(a.root) / ex, len(gen))], 2), crf=20)
+        print("wrote", a.out / f"{ex}_stream.mp4", flush=True)
+
+
+def release_configs(a):
+    from cwm_h3_inference.constants import STRIDE_FRAMES, WINDOW_FRAMES
+    windows = 1 + max(0, -(-(a.chunks * CLIP_FRAMES - WINDOW_FRAMES) // STRIDE_FRAMES))
+    for e in a.examples:
+        ex = Path(a.root) / e
+        cfg = json.loads((ex / "config.json").read_text())
+        d = (a.out / e).resolve()
+        d.mkdir(parents=True, exist_ok=True)
+        absolute = lambda v: str((ex / v).resolve())
+        cfg.update(anchor_image=absolute(cfg["anchor_image"]), condition_root=absolute(cfg["condition_root"]),
+                   models={k: absolute(v) for k, v in cfg["models"].items()}, windows=cfg["windows"][:windows],
+                   cache_dir=str(d / "cache"), output=str(d / "output.mp4"), blocks_to_swap=16)
+        (d / "config.json").write_text(json.dumps(cfg, indent=1))
+        print(f"{e}: {windows} windows, {WINDOW_FRAMES + STRIDE_FRAMES * (windows - 1)} frames -> {d / 'config.json'}")
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("stage", choices=["prepare", "generate", "release"])
+    ap.add_argument("--examples", nargs="+", required=True, help="prepare: example directories; generate: their names")
+    ap.add_argument("--root", default=str(Path.home() / "projects/code-world-model/release/examples"))
+    ap.add_argument("--chunks", type=int, default=24)
+    ap.add_argument("--lora", type=Path)
+    ap.add_argument("--out", type=Path)
+    ap.add_argument("--steps", type=int, default=35)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--fp8", action="store_true", help="the blocks' large linears in FP8 (Hopper)")
+    ap.add_argument("--rule", choices=["euler", "x0"], default="euler", help="x0: few-step (distilled) sampling")
+    ap.add_argument("--shift", type=float, default=0.0, help="sampling flow shift (0: the released 12)")
+    ap.add_argument("--kv-reuse", action="store_true", help="carry the history's keys/values across chunks")
+    ap.add_argument("--history-from-last-step", action=argparse.BooleanOptionalAction, default=True,
+                    help="kv-reuse: cache each chunk from its last denoising step (one forward fewer per chunk; "
+                         "--no-history-from-last-step: one more clean forward)")
+    ap.add_argument("--history-noise", type=float, default=0.0, help="noise level of the history chunks")
+    a = ap.parse_args()
+    if a.stage == "prepare":
+        prepare([Path(e) if Path(e).is_dir() else Path(a.root) / e for e in a.examples], a.chunks)
+    elif a.stage == "release":
+        release_configs(a)
+    else:
+        generate(a)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/eval/video.py b/rtcwm/eval/video.py
new file mode 100644
index 0000000000000000000000000000000000000000..f826352c108bf61fd7e11ae855489e084fdec386
--- /dev/null
+++ b/rtcwm/eval/video.py
@@ -0,0 +1,47 @@
+"""Encoding frames / decoding latents with the H3 video VAE, and writing / reading mp4s."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import av
+import numpy as np
+import torch
+
+from .. import paths
+
+
+def load_vae(device):
+    from musubi_tuner.minimax_h3.video_vae import load_video_vae
+    return load_video_vae(paths.VIDEO_VAE, device=device, dtype=torch.float16, disable_mmap=False).eval()
+
+
+@torch.inference_mode()
+def encode_frames(vae, frames: np.ndarray, device) -> torch.Tensor:
+    """[frames,768,1344,3] uint8 -> [1,24,F',48,84] latents, as the pre-encoder encodes targets / keyframes."""
+    from musubi_tuner.minimax_h3_generate_video import _prepare_pixels
+    px = _prepare_pixels(torch.from_numpy(np.ascontiguousarray(frames))).to(device, torch.float16)
+    return vae.encode(px).float().cpu()
+
+
+@torch.inference_mode()
+def decode(vae, latents, device) -> np.ndarray:
+    """[1,24,F,48,84] latents -> [frames,768,1344,3] uint8."""
+    px = vae.decode(latents.to(device, torch.float16))[0].float().cpu()            # [3,frames,H,W] in [-1,1]
+    return ((px.clamp(-1, 1) + 1) * 127.5).round().byte().permute(1, 2, 3, 0).numpy()
+
+
+def write_mp4(path: Path, frames: np.ndarray, fps: int = 24, crf: int = 20):
+    with av.open(str(path), "w") as c:
+        s = c.add_stream("libx264", rate=fps)
+        s.width, s.height, s.pix_fmt = frames.shape[2], frames.shape[1], "yuv420p"
+        s.options = {"crf": str(crf)}
+        for f in frames:
+            for p in s.encode(av.VideoFrame.from_ndarray(np.ascontiguousarray(f), format="rgb24")):
+                c.mux(p)
+        for p in s.encode():
+            c.mux(p)
+
+
+def read_mp4(path: Path) -> np.ndarray:
+    with av.open(str(path)) as c:
+        return np.stack([f.to_ndarray(format="rgb24") for f in c.decode(video=0)])
diff --git a/rtcwm/model/__init__.py b/rtcwm/model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5032162b08a0886ec81da39733e1eef0aed23f7c
--- /dev/null
+++ b/rtcwm/model/__init__.py
@@ -0,0 +1,16 @@
+"""Causal CWM: the released MiniMax-H3 Ref2VA model + CWM LoRA, run chunk-causally.
+
+A 124-frame CWM window (37 latents) is packed as
+    condition : text | Picture 1 anchor | Video 1 coarse-proxy DUV | keyframe (target latent 0)
+    audio     : clean silence (CARLA is silent)
+    video     : target latents 1..36 = 7 chunks of 6, 5, 5, 5, 5, 5, 5 latents (each generated prefix is 5n+2
+                latents, aligned with the VAE's 17-frame clips; see rtcwm.model.layout)
+Two training objectives (rtcwm.model.objectives):
+    df : diffusion forcing -- one copy of the video, every chunk at its own noise level; the condition,
+         audio and chunk 0 see each other, chunk c >= 1 sees the condition, audio and chunks c-5..c.
+         Condition rows have fixed (clean) times, so a sampler caches their keys/values together with
+         chunk 0's (rtcwm.model.sampler.KVSampler).
+    tf : teacher forcing (SolarWM Stage1 layout) -- clean and noisy copies of the video; noisy chunk c
+         sees the condition, audio, clean chunks c-5..c-1 and itself; condition times follow the
+         video time as in the released CWM.
+"""
diff --git a/rtcwm/model/attention.py b/rtcwm/model/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..2235be01ee4bc94d82ddfe3fc354e993400ea711
--- /dev/null
+++ b/rtcwm/model/attention.py
@@ -0,0 +1,181 @@
+"""Masked / cached attention for the MiniMax-H3 blocks.
+
+musubi's H3 Attention runs dense FlashAttention over the whole packed sequence. install() patches its
+forward with two extra modes, selected per block call through `active(...)`:
+
+  mask : torch flex_attention with a BlockMask over the whole packed sequence (training, reference)
+  dense: a masks.DenseDF instead of a BlockMask -- the same training mask as FlashAttention-3 prefix calls
+  kv   : the call's rows attend to cached keys/values of earlier rows plus their own (KV-cache sampling,
+         FlashAttention-3; see KV_ATTENTION);
+         with store=list the call's own keys/values are appended per layer. A mask may be combined with
+         kv to fill the cache (chunk 0 together with the condition rows).
+
+The mode is set inside each block call (not around the loop) so that activation-checkpoint
+recomputation in backward sees it too.
+"""
+from __future__ import annotations
+
+import contextlib
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import flex_attention
+
+from musubi_tuner.minimax_h3 import model as h3model
+
+from .layout import CHUNK_SIZES
+from .masks import CHUNK_ROWS, DenseDF
+
+_compiled_flex = torch.compile(flex_attention, dynamic=True)
+# KV-cache attention over the concatenated past: "fa3" FlashAttention-3 (23% faster per sampled chunk);
+# "sdpa" torch SDPA (the earlier behaviour, kept for tests)
+KV_ATTENTION = "fa3"
+_state = {"mask": None, "kv": None}
+_orig_forward = h3model.Attention.forward
+_orig_rope = h3model._apply_rope_split_half              # kept for tests/test_samplers.py
+
+
+def _forward(self, hidden_states, rotation_table=None):
+    mask, kv = _state["mask"], _state["kv"]
+    if mask is None and kv is None:
+        return _orig_forward(self, hidden_states, rotation_table)
+    b, s, _ = hidden_states.shape
+    q, k, v = self.qkv_proj(hidden_states).split(self.inner_dim, dim=-1)
+    q = self.q_norm(q.reshape(b, s, self.num_heads, self.head_dim))
+    k = self.k_norm(k.reshape(b, s, self.num_heads, self.head_dim))
+    v = v.reshape(b, s, self.num_heads, self.head_dim)
+    if rotation_table is not None:
+        q = rope(q, rotation_table)
+        k = rope(k, rotation_table)
+    dt = hidden_states.dtype if hidden_states.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
+    if isinstance(mask, DenseDF):
+        out = dense_df(q.to(dt), k.to(dt), v.to(dt), mask)
+        return self.out_proj(out.reshape(b, s, self.inner_dim))
+    q, k, v = (t.transpose(1, 2).contiguous().to(dt) for t in (q, k, v))     # [B,H,S,D]
+    if kv is not None and kv.get("store") is not None:
+        kv["store"].append((k, v))
+    if mask is not None:
+        out = _compiled_flex(q, k, v, block_mask=mask)
+    else:
+        pieces = kv["past"][self._rtcwm_layer] if kv.get("past") is not None else []
+        if kv.get("causal_blocks"):                  # the call's rows are blocks of n rows, block p seeing the
+            n = kv["causal_blocks"]                  # past and blocks <= p (the proxy rows of a span)
+            pk = torch.cat([_key(p) for p in pieces], 2).transpose(1, 2)
+            pv = torch.cat([p[1] for p in pieces], 2).transpose(1, 2)
+            qs, ks, vs = (t.transpose(1, 2) for t in (q, k, v))
+            out = torch.cat([_fa3(qs[:, i:i + n], torch.cat([pk, ks[:, :i + n]], 1), torch.cat([pv, vs[:, :i + n]], 1))
+                             for i in range(0, s, n)], 1)
+            return self.out_proj(out.reshape(b, s, self.inner_dim))
+        if pieces:                                   # concatenated per layer: one layer's copy at a time;
+            k = torch.cat([_key(p) for p in pieces] + [k], 2)   # a piece may carry a pending time shift
+            v = torch.cat([p[1] for p in pieces] + [v], 2)
+        if KV_ATTENTION == "sdpa":
+            out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2)
+        else:
+            out = _fa3(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2))
+        return self.out_proj(out.reshape(b, s, self.inner_dim))
+    return self.out_proj(out.transpose(1, 2).reshape(b, s, self.inner_dim))
+
+
+def _key(piece):
+    """Cached keys [B,H,S,D] of a piece (k, v) or (k, v, rotation): the rotation (rtcwm.model.sampler._moved)
+    moves them on the time axis -- RoPE rotations compose -- without storing a moved copy."""
+    if len(piece) == 2:
+        return piece[0]
+    return rope(piece[0].transpose(1, 2), piece[2]).transpose(1, 2)
+
+
+def _fa3(q, k, v):
+    from flash_attn_interface import flash_attn_func
+    out = flash_attn_func(q, k, v)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def dense_df(q, k, v, lay: DenseDF):
+    """Attention of a packed training span [head | DUV 0..5 | audio | video 0..n-1] under df_mask, as
+    unmasked FlashAttention-3 calls over key prefixes (see masks.DenseDF). q, k, v [B,S,H,D]."""
+    h, a = lay.head_rows, lay.audio_rows
+    dr, cr = lay.duv_rows // len(CHUNK_SIZES), CHUNK_ROWS
+    d1 = h + lay.duv_rows
+    v0 = d1 + a
+    assert q.shape[1] == v0 + lay.n_chunks * cr, (q.shape, lay)
+    duv = lambda t, p: t[:, h + p * dr:h + (p + 1) * dr]
+    vid = lambda t, c: t[:, v0 + c * cr:v0 + (c + 1) * cr]
+
+    def main(t):
+        parts = [t[:, :h], t[:, d1:v0]]
+        for c in range(lay.n_chunks):
+            parts += [duv(t, c), vid(t, c)]
+        return torch.cat(parts, 1)
+
+    km, vm = main(k), main(v)
+    kp, vp = (torch.cat([t[:, :h], t[:, d1:v0], t[:, h:d1]], 1) for t in (k, v))
+    st = h + a                                                   # static keys lead both orders
+    n = st + (dr + cr if lay.static_sees_first else 0)
+    o_st = _fa3(torch.cat([q[:, :h], q[:, d1:v0]], 1), km[:, :n], vm[:, :n])
+    o_duv = [_fa3(duv(q, p), kp[:, :st + (p + 1) * dr], vp[:, :st + (p + 1) * dr]) for p in range(len(CHUNK_SIZES))]
+    o_vid = [_fa3(vid(q, c), km[:, :st + (c + 1) * (dr + cr)], vm[:, :st + (c + 1) * (dr + cr)])
+             for c in range(lay.n_chunks)]
+    return torch.cat([o_st[:, :h], *o_duv, o_st[:, h:], *o_vid], 1)
+
+
+def rope(hidden_states, rotation_table):
+    """musubi's _apply_rope_split_half written element-wise. The original multiplies every (token, head,
+    pair) by its 2x2 rotation with a batched matmul -- one tiny cuBLAS gemv per few thousand elements,
+    ~40% of a sampling step. Same rotation table, accumulated in fp32 like the matmul."""
+    pairs = rotation_table.shape[-3]
+    r = rotation_table.float()                                  # [B,S,1,pairs,2,2] = [[cos, -sin], [sin, cos]]
+    a, b = hidden_states[..., :pairs].float(), hidden_states[..., pairs:2 * pairs].float()
+    rotated = torch.cat((r[..., 0, 0] * a + r[..., 0, 1] * b, r[..., 1, 0] * a + r[..., 1, 1] * b), dim=-1)
+    return torch.cat((rotated.to(hidden_states.dtype), hidden_states[..., 2 * pairs:]), dim=-1)
+
+
+_orig_mod_scale_shift, _orig_mod_gate = h3model._mod_scale_shift, h3model._mod_gate   # kept for tests
+_row_index = {}
+
+
+def _rows(batch_segments, device):
+    """Per-row modulation index of the (single-sample) segment plan, built once per forward: every block
+    of a forward call shares the same plan object."""
+    key = (id(batch_segments), device)
+    hit = _row_index.get(key)
+    if hit is not None and hit[0] is batch_segments:
+        return hit[1]
+    (segments,) = batch_segments
+    seg = torch.tensor(segments, dtype=torch.long)
+    idx = torch.repeat_interleave(seg[:, 2], seg[:, 1] - seg[:, 0]).to(device, non_blocking=True)
+    _row_index.clear()
+    _row_index[key] = (batch_segments, idx)
+    return idx
+
+
+def mod_scale_shift(hidden_states, shift, scale, batch_segments):
+    """musubi's _mod_scale_shift without the Python loop over segments (hundreds of small kernels per block
+    with the text's token tags): the same x * (1 + scale) + shift, gathered per row."""
+    idx = _rows(batch_segments, hidden_states.device)
+    return hidden_states * (1.0 + scale[idx]).unsqueeze(0) + shift[idx].unsqueeze(0)
+
+
+def mod_gate(residual, update, gate, batch_segments):
+    """musubi's _mod_gate, vectorised the same way: residual + update * gate."""
+    idx = _rows(batch_segments, residual.device)
+    return residual + update * gate[idx].unsqueeze(0)
+
+
+def install(model) -> None:
+    """Patch H3 attention, its RoPE and its AdaLN modulation (idempotent) and number the layers for the KV
+    cache."""
+    h3model.Attention.forward = _forward
+    h3model._apply_rope_split_half = rope                      # the released (unmasked) forward too
+    h3model._mod_scale_shift, h3model._mod_gate = mod_scale_shift, mod_gate
+    for i, blk in enumerate(model.blocks):
+        (blk.attn if hasattr(blk, "attn") else blk.module.attn)._rtcwm_layer = i    # FSDP-wrapped or not
+
+
+@contextlib.contextmanager
+def active(mask=None, kv=None):
+    _state["mask"], _state["kv"] = mask, kv
+    try:
+        yield
+    finally:
+        _state["mask"], _state["kv"] = None, None
diff --git a/rtcwm/model/fp8.py b/rtcwm/model/fp8.py
new file mode 100644
index 0000000000000000000000000000000000000000..129f34959f459c7c5eaa923dc3d86d6d8c86a851
--- /dev/null
+++ b/rtcwm/model/fp8.py
@@ -0,0 +1,43 @@
+"""FP8 linears for sampling (Hopper): the large bias-free linears of every DiT block -- attention qkv / out
+and the MLP -- with e4m3 weights scaled per output channel and activations quantized per row on the fly,
+multiplied by torch._scaled_mm into bf16. Inference only (no gradient), on a model whose LoRA is merged;
+approximate (FP8 rounding), so its effect is measured, not assumed (tests/test_fp8.py).
+"""
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+E4M3_MAX = 448.0
+
+
+class FP8Linear(nn.Module):
+    def __init__(self, linear: nn.Linear):
+        super().__init__()
+        assert linear.bias is None
+        w = linear.weight.detach().float()                                     # [out, in]
+        scale = (w.abs().amax(dim=1, keepdim=True) / E4M3_MAX).clamp_min(1e-12)
+        self.register_buffer("weight_fp8", (w / scale).to(torch.float8_e4m3fn), persistent=False)
+        self.register_buffer("weight_scale", scale.t().contiguous(), persistent=False)    # [1, out]
+        self.out_features = linear.out_features
+
+    def forward(self, x):
+        shape = x.shape
+        x2 = x.reshape(-1, shape[-1])
+        scale = (x2.abs().amax(dim=1, keepdim=True).float() / E4M3_MAX).clamp_min(1e-12)   # [rows, 1]
+        xq = (x2.float() / scale).to(torch.float8_e4m3fn)
+        y = torch._scaled_mm(xq, self.weight_fp8.t(), scale_a=scale, scale_b=self.weight_scale,
+                             out_dtype=torch.bfloat16)
+        return y.reshape(*shape[:-1], self.out_features)
+
+
+def quantize_blocks(model) -> int:
+    """Swap the DiT blocks' qkv / out / fc1 / fc2 linears for FP8Linear in place; returns how many."""
+    n = 0
+    for block in model.blocks:
+        for parent, names in ((block.attn, ("qkv_proj", "out_proj")), (block.mlp, ("fc1", "fc2"))):
+            for name in names:
+                setattr(parent, name, FP8Linear(getattr(parent, name)))
+                n += 1
+    torch.cuda.empty_cache()
+    return n
diff --git a/rtcwm/model/layout.py b/rtcwm/model/layout.py
new file mode 100644
index 0000000000000000000000000000000000000000..663c2229afd91c07bcdfb817858857f7999e44ad
--- /dev/null
+++ b/rtcwm/model/layout.py
@@ -0,0 +1,51 @@
+"""Streaming chunk layout and the per-row bookkeeping shared by the objective and the sampler.
+
+A clip is one continuous latent stream (rtcwm.train.preencode --stage stream). The H3 video VAE encodes
+independent 17-frame clips into 5 latents each; one chunk is one such clip (17 frames, 0.71 s), so a chunk
+never splits a VAE clip and every generated prefix ends on a clip boundary. There is no keyframe: chunk k
+is generated from the conditions and the chunks before it in its attention span, which slides with k
+(there is no generation window and no restart anywhere in the stream). The span (default 6 chunks: 30
+latents, 4.25 s; RTCWM_SPAN_CHUNKS sets it, for training and sampling alike -- a LoRA records the span it
+was trained with) is packed as a standalone H3 sequence with positions counted from the span's first chunk
+(rtcwm.model.packing.Window), so a late chunk never leaves the position range of the released model.
+"""
+from __future__ import annotations
+
+import os
+from itertools import accumulate
+
+CHUNK_LATENTS = 5                    # one 17-frame VAE clip
+CLIP_FRAMES = 17
+SPAN_CHUNKS = int(os.environ.get("RTCWM_SPAN_CHUNKS", "6"))   # attention span: a chunk sees itself and the ones before it
+CHUNK_SIZES = (CHUNK_LATENTS,) * SPAN_CHUNKS
+NUM_CHUNKS = SPAN_CHUNKS             # chunks in one packed span
+TARGET_LATENTS = sum(CHUNK_SIZES)    # video latents per packed span (30 for 6 chunks)
+CHUNK_STARTS = tuple(accumulate((0,) + CHUNK_SIZES[:-1]))
+FRAME_ROWS = 1008                    # (48/2) * (84/2) packed rows per latent frame
+LATENT_SHAPE = (24, 48, 84)          # channels, height, width
+
+# condition row times (model time: 1 = clean)
+TEXT_T_FIXED, VISUAL_T, CLEAN_T = 1.0, 0.999, 1.0
+
+
+def chunk_latents(c: int) -> slice:
+    """Latents of chunk c of a packed span."""
+    return slice(CHUNK_STARTS[c], CHUNK_STARTS[c] + CHUNK_SIZES[c])
+
+
+def chunk_rows(c: int) -> int:
+    return CHUNK_SIZES[c] * FRAME_ROWS
+
+
+def chunk_row_slice(c: int) -> slice:
+    """Packed rows of chunk c within the span's video rows."""
+    return slice(CHUNK_STARTS[c] * FRAME_ROWS, (CHUNK_STARTS[c] + CHUNK_SIZES[c]) * FRAME_ROWS)
+
+
+def chunk_shape(c: int = 0) -> tuple:
+    return (1, LATENT_SHAPE[0], CHUNK_SIZES[c], *LATENT_SHAPE[1:])
+
+
+def span_start(k: int) -> int:
+    """First chunk of the attention span of stream chunk k."""
+    return max(0, k - (SPAN_CHUNKS - 1))
diff --git a/rtcwm/model/loading.py b/rtcwm/model/loading.py
new file mode 100644
index 0000000000000000000000000000000000000000..8678ef93d436d1af27dac8874c92d60537c3c491
--- /dev/null
+++ b/rtcwm/model/loading.py
@@ -0,0 +1,157 @@
+"""Loading the frozen H3 transformer with a LoRA (trainable or merged), and sharding its blocks."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from .. import paths
+from . import attention
+
+EXPECTED_LORA_MODULES = 200            # the CWM LoRA: qkv / out / fc1 / fc2 of the 50 blocks
+
+
+def load_base():
+    """The released Ref2VA H3 transformer on CPU, bf16, frozen."""
+    from musubi_tuner.minimax_h3.model import load_h3_transformer
+    tr = load_h3_transformer(paths.DIT, device="cpu", dtype=torch.bfloat16, attn_mode="flash3", split_attn=False,
+                             disable_mmap=False)
+    tr.requires_grad_(False)
+    return tr
+
+
+def attach_trainable_lora(tr, lora: Path):
+    """Wrap the blocks' linears with a LoRA initialised from `lora`. Returns the (fp32-ready) network."""
+    from musubi_tuner.networks import lora_minimax_h3
+    sd = load_file(str(lora))
+    net = lora_minimax_h3.create_arch_network_from_weights(1.0, sd, unet=tr, for_inference=False)
+    net.apply_to(None, tr, apply_text_encoder=False, apply_unet=True)
+    info = net.load_state_dict(sd, strict=False)
+    assert not info.unexpected_keys, info.unexpected_keys[:5]
+    assert not info.missing_keys and len(net.unet_loras) == EXPECTED_LORA_MODULES, (len(info.missing_keys), len(net.unet_loras))
+    return net
+
+
+def _skip_when_off(m):
+    full = m.forward
+
+    def forward(x):
+        return m.org_forward(x) if m.multiplier == 0 else full(x)
+    return forward
+
+
+class LoRAStack:
+    """Several LoRAs on one frozen base (distillation: student, critic, teacher). They are stacked on
+    the same linears; exactly one is active at a time (the others get multiplier 0). Activation
+    checkpointing recomputes blocks in backward, so the network that built a graph must be active
+    again when that graph is back-propagated."""
+
+    def __init__(self, tr, specs: dict):
+        """specs: name -> (initial LoRA file, trainable)."""
+        self.nets, self.trainable = {}, {}
+        for name, (path, trainable) in specs.items():
+            net = attach_trainable_lora(tr, path)
+            net.requires_grad_(trainable)
+            self.nets[name], self.trainable[name] = net, trainable
+            for m in net.unet_loras:                      # musubi's LoRA computes its branch even at
+                m.forward = _skip_when_off(m)             # multiplier 0: every forward paid for all three
+        self.active = None
+
+    def to(self, device):
+        for net in self.nets.values():
+            net.to(device, torch.float32)
+        return self
+
+    def use(self, name: str):
+        """Activate one LoRA; only it (if trainable) requires gradients, so a backward pass produces
+        (and all-reduces) gradients for that network alone."""
+        for n, net in self.nets.items():
+            for m in net.unet_loras:
+                m.multiplier = 1.0 if n == name else 0.0
+            net.requires_grad_(n == name and self.trainable[n])
+        self.active = name
+
+    def params(self, name: str):
+        return list(self.nets[name].parameters())
+
+
+def load_for_inference(lora: Path, device, fp8: bool = False):
+    """Base + merged LoRA on `device`, eval mode, causal attention installed; fp8: the blocks' large linears
+    in FP8 (rtcwm.model.fp8)."""
+    from cwm_h3_inference.runtime import strict_merge_lora
+    tr = load_base()
+    strict_merge_lora(tr, lora, EXPECTED_LORA_MODULES)
+    tr = tr.to(device).eval()
+    if fp8:
+        from .fp8 import quantize_blocks
+        quantize_blocks(tr)
+    attention.install(tr)
+    return tr
+
+
+def shard_blocks(tr, group, device):
+    """FSDP-shard each block over `group` (the GPUs of one node); everything else is small and replicated.
+    The LoRA weights live in the network module, outside FSDP, and stay replicated."""
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    from torch.distributed.fsdp import ShardingStrategy
+    for name, child in tr.named_children():
+        if name != "blocks":
+            child.to(device)
+    from torch.distributed.fsdp._runtime_utils import _reshard, _unshard
+
+    # FSDP (v1) leaves a block's full parameters gathered after a forward run without autograd and after
+    # a forward inside our activation checkpoint -- every block then stays gathered (+55-60 GB per GPU
+    # during a KV-cache rollout, a no-grad teacher pass or a student gradient step). And as the base
+    # weights are frozen, FSDP registers no post-backward hook either, so a block gathered for its
+    # backward stays gathered as well. We therefore manage the block's gather/reshard ourselves:
+    def in_backward():
+        return torch._C._current_graph_task_id() != -1
+
+    def gather_for_recompute(module, _inputs):
+        # FSDP gathers a block for backward once per backward pass; a block used by several gradient
+        # forwards (several student chunks) is recomputed again after we resharded it
+        if in_backward() and module._handle.needs_unshard():
+            _unshard(module, module._handle, module._unshard_stream, module._pre_unshard_stream)
+            torch.cuda.current_stream().wait_stream(module._unshard_stream)
+
+    def reshard_after_forward(module, inputs, _output):
+        if in_backward():                  # checkpoint recomputation: the backward of the block follows
+            return
+        _reshard(module, module._handle, True)
+        h = inputs[0] if inputs and isinstance(inputs[0], torch.Tensor) else None
+        if h is not None and h.requires_grad:             # the block's backward is done once its input has a gradient
+            h.register_hook(lambda g: (_reshard(module, module._handle, True), g)[1])
+
+    # Every block is its own FSDP root, so FSDP's forward_prefetch (which follows one root's execution
+    # order) never gathers the next block. Outside backward we do it by hand: once block i holds its
+    # parameters (a pre-hook on the wrapped module, i.e. inside FSDP's forward, after its unshard), block
+    # i+1's all-gather is issued on its own unshard stream and overlaps block i's compute; FSDP's
+    # pre-forward of block i+1 then finds it prefetched and only waits for that stream. (The same in
+    # backward -- block i-1 during block i's recomputation -- measured no faster: the backward is compute
+    # bound.)
+    from torch.distributed.fsdp._flat_param import HandleTrainingState
+
+    def prefetch_next(i):
+        def hook(_module, _inputs):
+            if in_backward() or i + 1 >= len(tr.blocks):
+                return
+            nxt = tr.blocks[i + 1]
+            h = nxt._handle
+            if (h is None or getattr(nxt, "_unshard_stream", None) is None       # not run yet: FSDP initialises
+                    or not hasattr(h.flat_param, "_full_param_padded")            # a block lazily, in its first
+                    or h._prefetched or not h.needs_unshard()):                  # forward
+                return
+            prev, h._training_state = h._training_state, HandleTrainingState.FORWARD
+            _unshard(nxt, h, nxt._unshard_stream, nxt._pre_unshard_stream)
+            h._training_state, h._prefetched = prev, True
+        return hook
+
+    for i in range(len(tr.blocks)):
+        tr.blocks[i] = FSDP(tr.blocks[i], process_group=group, device_id=device, use_orig_params=True,
+                            sharding_strategy=ShardingStrategy.FULL_SHARD, limit_all_gathers=True)
+        tr.blocks[i].register_forward_pre_hook(gather_for_recompute)
+        tr.blocks[i].register_forward_hook(reshard_after_forward)
+        tr.blocks[i].module.register_forward_pre_hook(prefetch_next(i))
+    attention.install(tr)
+    return tr
diff --git a/rtcwm/model/masks.py b/rtcwm/model/masks.py
new file mode 100644
index 0000000000000000000000000000000000000000..96e5b6e886b49276bf046259bd25f090670e6760
--- /dev/null
+++ b/rtcwm/model/masks.py
@@ -0,0 +1,132 @@
+"""Block-causal flex-attention mask over a packed span.
+
+Rows come in contiguous runs of one (role, chunk). The BlockMask's full / partial 128x128 tiles are
+derived from those few runs (never a token-level [S,S] mask; the builder follows SolarWM's
+backends/minimax_h3/torch_mask.py, Apache-2.0), and mask_mod decides single tokens from a handful of
+captured scalars, so the compiled kernel does not specialise on the caption length.
+
+A span is packed [text | Picture 1 | DUV chunks 0..5 | audio | video chunks first .. first+n-1]. Roles:
+STATIC (text, Picture 1, audio), DUV (the coarse proxy, one run per chunk) and VIDEO. Nothing sees a
+proxy chunk later than its own (real-time control: the future proxy does not exist at inference).
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch.nn.attention.flex_attention import BlockMask
+
+from .layout import CHUNK_SIZES, FRAME_ROWS, NUM_CHUNKS, SPAN_CHUNKS
+
+STATIC, DUV, VIDEO = 0, 1, 2
+TILE = 128
+assert len(set(CHUNK_SIZES)) == 1
+CHUNK_ROWS = CHUNK_SIZES[0] * FRAME_ROWS
+
+
+def df_allows(qr, qc, kr, kc, static_sees_first: bool) -> bool:
+    """Diffusion forcing with a causal proxy.
+        static    : static rows, and -- static_sees_first, i.e. the span starts the stream -- proxy and video
+                    chunk 0 (with no history, chunk 0 needs condition rows that see it, as in the released
+                    model; with a history, conditions that see only themselves do better)
+        proxy p   : static rows and proxy chunks <= p
+        video c   : static rows, proxy chunks <= c and video chunks c-5..c"""
+    if qr == STATIC:
+        return kr == STATIC or (static_sees_first and kc == 0 and kr in (DUV, VIDEO))
+    if qr == DUV:
+        return kr == STATIC or (kr == DUV and kc <= qc)
+    return kr == STATIC or (kr == DUV and kc <= qc) or (kr == VIDEO and 0 <= qc - kc < SPAN_CHUNKS)
+
+
+def _tiles(runs, allows, S, device):
+    allowed = {(r, c): [(a, b) for a, b, r2, c2 in runs if allows(r, c, r2, c2)] for _, _, r, c in runs}
+    nb = math.ceil(S / TILE)
+    ks = torch.arange(nb) * TILE
+    ke = torch.clamp(ks + TILE, max=S)
+    full = torch.zeros(nb, nb, dtype=torch.bool)
+    anyt = torch.zeros(nb, nb, dtype=torch.bool)
+    for qb in range(nb):
+        q0, q1 = qb * TILE, min(qb * TILE + TILE, S)
+        rf, ra = torch.ones(nb, dtype=torch.bool), torch.zeros(nb, dtype=torch.bool)
+        for a, b, r, c in runs:
+            if a < q1 and q0 < b:
+                sf, sa = torch.zeros(nb, dtype=torch.bool), torch.zeros(nb, dtype=torch.bool)
+                for a2, b2 in allowed[(r, c)]:
+                    sf |= (a2 <= ks) & (ke <= b2)
+                    sa |= (ks < b2) & (a2 < ke)
+                rf &= sf
+                ra |= sa
+        full[qb], anyt[qb] = rf, ra
+    mb = 1 << (nb - 1).bit_length()
+
+    def meta(m):
+        o = torch.zeros(mb, mb, dtype=torch.int32)
+        o[:nb, :nb] = m
+        return (o.sum(-1).to(torch.int32)[None, None].to(device),
+                torch.argsort(o, dim=-1, descending=True, stable=True).to(torch.int32)[None, None].to(device))
+
+    pc, pi = meta(anyt & ~full)
+    fc, fi = meta(full)
+    return pc, pi, fc, fi
+
+
+def _scalars(device, *values):
+    return [torch.tensor(v, dtype=torch.long, device=device) for v in values]
+
+
+def span_runs(head_rows: int, duv_rows: int, audio_rows: int, n_chunks: int, first_chunk: int = 0):
+    """(runs, total rows) of [head | DUV x6 | audio | video chunks first..]; runs are (start, end, role, chunk)."""
+    assert duv_rows % NUM_CHUNKS == 0, duv_rows
+    dr = duv_rows // NUM_CHUNKS
+    runs = [(0, head_rows, STATIC, -1)]
+    runs += [(head_rows + p * dr, head_rows + (p + 1) * dr, DUV, p) for p in range(NUM_CHUNKS)]
+    v0 = head_rows + duv_rows + audio_rows
+    runs.append((head_rows + duv_rows, v0, STATIC, -1))
+    runs += [(v0 + j * CHUNK_ROWS, v0 + (j + 1) * CHUNK_ROWS, VIDEO, first_chunk + j) for j in range(n_chunks)]
+    return [r for r in runs if r[1] > r[0]], v0 + n_chunks * CHUNK_ROWS
+
+
+def df_mask(head_rows: int, duv_rows: int, audio_rows: int, n_chunks: int, device, first_chunk: int = 0,
+            static_sees_first: bool = True) -> BlockMask:
+    """Mask of a packed span [text | Picture 1 | DUV chunks 0..5 | audio | video chunks first ..]."""
+    runs, S = span_runs(head_rows, duv_rows, audio_rows, n_chunks, first_chunk)
+    pc, pi, fc, fi = _tiles(runs, lambda *a: df_allows(*a, static_sees_first), S, device)
+    d0, d1, v0, se, dr, cr, wb, first, ssf = _scalars(
+        device, head_rows, head_rows + duv_rows, head_rows + duv_rows + audio_rows, S, duv_rows // NUM_CHUNKS,
+        CHUNK_ROWS, SPAN_CHUNKS, first_chunk, int(static_sees_first))
+
+    def classify(i):
+        duv = (i >= d0) & (i < d1)
+        vid = (i >= v0) & (i < se)
+        chunk = torch.where(duv, torch.div(i - d0, dr, rounding_mode="floor"),
+                            torch.div(i - v0, cr, rounding_mode="floor") + first)
+        return ~duv & ~vid, duv, vid, chunk
+
+    def mask_mod(_b, _h, qi, ki):
+        q_st, q_du, q_vi, qc = classify(qi)
+        k_st, k_du, k_vi, kc = classify(ki)
+        first_rows = (ssf == 1) & (kc == 0) & (k_du | k_vi)
+        return ((q_st & (k_st | first_rows)) | (q_du & (k_st | (k_du & (kc <= qc))))
+                | (q_vi & (k_st | (k_du & (kc <= qc)) | (k_vi & (qc - kc >= 0) & (qc - kc < wb)))))
+
+    return BlockMask.from_kv_blocks(pc, pi, fc, fi, BLOCK_SIZE=TILE, mask_mod=mask_mod, seq_lengths=(S, S))
+
+
+@dataclass(frozen=True)
+class DenseDF:
+    """The df_mask of a training span (video chunks 0..n-1, n <= SPAN_CHUNKS) for dense FlashAttention.
+
+    Under the rule every query run's keys are one contiguous prefix of one of two key orders:
+        main  [static | DUV 0 | video 0 | DUV 1 | video 1 | ...]   static rows, video chunk c
+        proxy [static | DUV 0 | DUV 1 | ...]                        proxy chunk p
+    (static = text, Picture 1 and audio), so the masked attention is a few unmasked prefix calls with the
+    same result as the BlockMask, without its per-tile mask evaluation (rtcwm.model.attention)."""
+    head_rows: int
+    duv_rows: int
+    audio_rows: int
+    n_chunks: int
+    static_sees_first: bool = True
+
+    def __post_init__(self):
+        assert 0 < self.n_chunks <= SPAN_CHUNKS and self.duv_rows % NUM_CHUNKS == 0, self
diff --git a/rtcwm/model/objectives.py b/rtcwm/model/objectives.py
new file mode 100644
index 0000000000000000000000000000000000000000..f334266ee0ccbe15e3a7e07ecf5081e0b038186b
--- /dev/null
+++ b/rtcwm/model/objectives.py
@@ -0,0 +1,90 @@
+"""The causal training objective (diffusion forcing) on one packed window.
+
+CWM's flow-matching parameterisation: noisy = (1 - sigma) x + sigma n, model time 1 - sigma, target
+velocity x - n, sigma flow-shifted with CWM's video shift.
+"""
+from __future__ import annotations
+
+import torch
+
+from cwm_h3_inference.constants import VIDEO_SHIFT, VISUAL_CONDITION_CLEAN
+
+from .layout import CHUNK_SIZES, FRAME_ROWS, NUM_CHUNKS, TARGET_LATENTS, chunk_rows
+from .masks import DenseDF, df_mask
+from .packing import Window, run_blocks, video_output
+
+OBJECTIVES = ("df",)
+
+
+def shift(base, s: float = VIDEO_SHIFT):
+    """Flow-shifted noise amount, as in musubi's H3 trainer: s*t / (1 + (s-1)*t)."""
+    return s * base / (1.0 + (s - 1.0) * base)
+
+
+def _cached(cache: dict, key, build):
+    if key not in cache:
+        cache.clear()
+        cache[key] = build()
+    return cache[key]
+
+
+def df_forward(model, win: Window, video, chunk_t, mask_cache: dict, train: bool = True, offload: bool = False,
+               attn: str = "fa3"):
+    """video [1,24,30,H,W], each chunk noised to its level; chunk_t: 6 model times (1 = clean).
+    attn: fa3 (dense prefix calls) or flex (BlockMask), the same mask. Returns the predicted velocity of
+    all 30 latents."""
+    hidden = torch.cat([win.rows, win.video_rows(model, video)], 1)
+    rows = torch.tensor([chunk_rows(c) for c in range(NUM_CHUNKS)])
+    t_rows = torch.cat([win.cond_times(), torch.tensor(chunk_t, dtype=torch.float32).repeat_interleave(rows)])
+    tags = torch.cat([win.tags, torch.zeros(TARGET_LATENTS * FRAME_ROWS, dtype=torch.int64)])
+    if attn == "fa3":
+        mask = DenseDF(win.head_rows, win.duv_rows, win.audio_rows, NUM_CHUNKS, win.static_sees_first)
+    else:
+        key = ("df", win.head_rows, win.duv_rows, win.audio_rows, win.static_sees_first)
+        mask = _cached(mask_cache, key, lambda: df_mask(win.head_rows, win.duv_rows, win.audio_rows, NUM_CHUNKS, win.dev,
+                                                        static_sees_first=win.static_sees_first))
+    hidden, temb, inv = run_blocks(model, hidden, t_rows, tags, torch.cat([win.ca_pos, win.target_pos]),
+                                   mask=mask, train=train, offload=offload)
+    return video_output(model, hidden, temb, inv, win.rows.shape[1], TARGET_LATENTS)
+
+
+def augment_conditions(sample, dev, gen):
+    """Picture 1 and the DUV reference mixed 0.999 clean / 0.001 noise, as at CWM inference."""
+    ref, anc = sample["ref"].to(dev), sample["anchor"].to(dev)
+    ref = VISUAL_CONDITION_CLEAN * ref + (1 - VISUAL_CONDITION_CLEAN) * torch.randn(ref.shape, generator=gen).to(dev)
+    anc = VISUAL_CONDITION_CLEAN * anc + (1 - VISUAL_CONDITION_CLEAN) * torch.randn(anc.shape, generator=gen).to(dev)
+    return anc, ref
+
+
+def make_window(model, sample, dev, gen) -> Window:
+    anc, ref = augment_conditions(sample, dev, gen)
+    return Window(model, text_hidden=sample["hidden"], text_tags=sample["tags"], anchor=anc, duv=ref,
+                  audio=sample["audio"].to(dev), n_latents=sample["ref"].shape[2],
+                  static_sees_first=bool(sample.get("stream_start", True)))
+
+
+def df_levels(gen, level_shift: float = 3.0) -> torch.Tensor:
+    """Per-chunk noise amounts, independent per chunk and flow-shifted by level_shift. History chunks are never
+    the clean ground truth: a clean ground-truth history teaches the model to trust its history completely,
+    which it must not do with its own (imperfect) history at inference (v3 drew one with probability 0.5 and
+    drifted more); at a moderate shift the history comes at every noise level and the model learns to use it
+    with care."""
+    return shift(torch.rand(NUM_CHUNKS, generator=gen), level_shift)
+
+
+def df_loss(model, sample, dev, gen, mask_cache, offload: bool = False, attn: str = "fa3",
+            level_shift: float = 3.0):
+    x = sample["x"].to(dev)
+    sig = df_levels(gen, level_shift)
+    s_lat = sig.repeat_interleave(torch.tensor(CHUNK_SIZES)).to(dev).view(1, 1, -1, 1, 1)
+    n = torch.randn(x.shape, generator=gen).to(dev)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        win = make_window(model, sample, dev, gen)
+        pred = df_forward(model, win, (1 - s_lat) * x + s_lat * n, [float(1.0 - v) for v in sig], mask_cache,
+                          offload=offload, attn=attn)
+    noisy = (s_lat > 0).float()
+    se = ((pred.float() - (x - n).float()) ** 2) * noisy
+    return se.sum() / (noisy.sum() * x.shape[1] * x.shape[3] * x.shape[4]).clamp_min(1.0)
+
+
+LOSSES = dict(df=df_loss)
diff --git a/rtcwm/model/packing.py b/rtcwm/model/packing.py
new file mode 100644
index 0000000000000000000000000000000000000000..894c870da5dcdeadd1225c288b9eeb2d289a9620
--- /dev/null
+++ b/rtcwm/model/packing.py
@@ -0,0 +1,116 @@
+"""Packing one window's condition and audio rows, and running the H3 blocks on a packed sequence."""
+from __future__ import annotations
+
+import torch
+
+from musubi_tuner.minimax_h3.packing import (H3VideoGeometry, _audio_grid, _build_modulation_segments, _frame_grid,
+                                             _unpack_video_rows, _video_grid, _video_t_spans, pack_audio_rows,
+                                             pack_video_rows)
+
+from musubi_tuner.utils.model_utils import create_cpu_offloading_wrapper
+
+from . import attention
+from .layout import CLEAN_T, CLIP_FRAMES, FRAME_ROWS, LATENT_SHAPE, TARGET_LATENTS, TEXT_T_FIXED, VISUAL_T
+
+
+def audio_frames(n_latents: int) -> int:
+    """Audio latent frames that cover n_latents video latents (whole VAE clips), as the released layout counts them."""
+    return (10 * (n_latents // 5) * CLIP_FRAMES + 3) // 6
+
+
+def window_positions(text_len: int, anchor_hw, duv_geometry: H3VideoGeometry, n_audio: int, n_latents: int):
+    """(text, Picture 1, DUV, audio, video) positions of a window, laid out exactly as musubi's
+    build_position_grid lays out a Ref2VA window (tests/test_positions.py), for any number of latents:
+    one time axis -- text, Picture 1, the DUV video, then the audio and the target video."""
+    text = torch.zeros(text_len, 3, dtype=torch.float64)
+    text[:, 0] = torch.arange(text_len, dtype=torch.float64)
+    cursor = float(text_len)
+    frame, _ = _frame_grid(H3VideoGeometry(1, *anchor_hw))
+    image = torch.empty(frame.shape[0], 3, dtype=torch.float64)
+    image[:, 0], image[:, 1:] = cursor, frame
+    cursor += 1.0
+    duv = _video_grid(duv_geometry, cursor)
+    cursor += sum(_video_t_spans(duv_geometry.frames))
+    target = H3VideoGeometry(n_latents, *LATENT_SHAPE[1:])
+    _, width = _frame_grid(target)
+    audio = _audio_grid(cursor, n_audio, float(width[0]), float(width[-1]))
+    return text, image, duv, audio, _video_grid(target, cursor)
+
+
+class Window:
+    """Projected [text | Picture 1 | DUV | audio] rows of one span and the positions of its video latents.
+
+    Condition times are fixed (text 1, Picture 1 / DUV 0.999, audio 1), so the condition rows' keys and
+    values do not depend on the denoising step. static_sees_first: the span starts the stream, and its
+    static rows see chunk 0 (rtcwm.model.masks)."""
+
+    def __init__(self, model, *, text_hidden, text_tags, anchor, duv, audio, n_latents: int = TARGET_LATENTS,
+                 static_sees_first: bool = True):
+        dev = model.video_patch_proj.weight.device
+        self.L = text_hidden.shape[1]
+        text_pos, image_pos, duv_pos, audio_pos, video_pos = window_positions(
+            self.L, anchor.shape[3:], H3VideoGeometry(*duv.shape[2:]), audio.shape[-1], n_latents)
+        self.ca_pos = torch.cat([text_pos, image_pos, duv_pos, audio_pos])
+        self.target_pos = video_pos
+        vis = [pack_video_rows(t.to(dev)) for t in (anchor, duv)]
+        self.n_visual = vis[0].shape[1] + vis[1].shape[1]          # Picture 1 + DUV rows
+        self.head_rows, self.duv_rows = self.L + vis[0].shape[1], vis[1].shape[1]
+        self.static_sees_first = static_sees_first
+        assert image_pos.shape[0] == vis[0].shape[1] and duv_pos.shape[0] == vis[1].shape[1]
+        pv = model.video_patch_proj(torch.cat(vis, 1).to(torch.float32)).to(model.dtype)
+        pa = model.audio_patch_proj(pack_audio_rows(audio.to(dev)).to(torch.float32)).to(model.dtype)
+        text = model.token_refiner(model.condition_proj(text_hidden.to(dev, model.dtype)))
+        self.rows = torch.cat([text, pv, pa], 1)
+        self.cond_rows, self.audio_rows = self.L + pv.shape[1], pa.shape[1]
+        assert self.rows.shape[1] == self.ca_pos.shape[0]
+        tags = torch.zeros(self.rows.shape[1], dtype=torch.int64)
+        tags[:self.L] = text_tags[0].cpu()
+        tags[self.cond_rows:] = 2
+        self.tags, self.dev = tags, dev
+
+    def cond_times(self) -> torch.Tensor:
+        t = torch.full((self.rows.shape[1],), CLEAN_T)
+        t[:self.L] = TEXT_T_FIXED
+        t[self.L:self.L + self.n_visual] = VISUAL_T
+        return t
+
+    def video_rows(self, model, latents):
+        return model.video_patch_proj(pack_video_rows(latents.to(self.dev)).to(torch.float32)).to(model.dtype)
+
+
+def run_blocks(model, hidden, t_rows, tags, positions, *, mask=None, kv=None, train: bool = False,
+               offload: bool = False):
+    """All H3 blocks over a packed sequence with per-row times. Returns (hidden, time embeddings, row ->
+    unique-time index). train: activation checkpointing per block; offload: checkpointed block inputs kept on
+    the CPU."""
+    uniq, inv = torch.unique(t_rows, sorted=True, return_inverse=True)
+    segments = (_build_modulation_segments(3 * inv + tags),)
+    temb = model.time_embedder(uniq.to(hidden.device)).to(model.dtype)
+    rot = model._rotation_table(positions.unsqueeze(0).to(hidden.device), hidden.dtype)
+
+    def wrap(block):
+        def run(*args):
+            with attention.active(mask, kv):
+                return block(*args)
+        return run
+
+    dev = hidden.device
+    for block in model.blocks:
+        fn = wrap(block)
+        if train:
+            if offload:
+                fn = create_cpu_offloading_wrapper(fn, dev)
+            hidden = torch.utils.checkpoint.checkpoint(fn, hidden, temb, segments, rot, use_reentrant=False)
+        else:
+            hidden = fn(hidden, temb, segments, rot)
+    return hidden.to(dev), temb, inv          # the offloading wrapper returns the last output on the CPU
+
+
+def video_output(model, hidden, temb, inv, start: int, n_latents: int):
+    """Final layer on rows [start, start + n_latents frames) -> velocity [1,24,n_latents,48,84]."""
+    rows = slice(start, start + n_latents * FRAME_ROWS)
+    empty = slice(start, start)
+    vid, _ = model.final_layer(hidden, temb, video_slice=rows, audio_slice=empty,
+                               video_timestep_index=int(inv[start]), audio_timestep_index=int(inv[start]),
+                               video_timestep_indices=inv[rows].unsqueeze(0), audio_timestep_indices=inv[empty].unsqueeze(0))
+    return _unpack_video_rows(vid, H3VideoGeometry(n_latents, *LATENT_SHAPE[1:]))
diff --git a/rtcwm/model/sampler.py b/rtcwm/model/sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9606fbc3c810b751b3b5dfaffd2f0f21284f52d
--- /dev/null
+++ b/rtcwm/model/sampler.py
@@ -0,0 +1,299 @@
+"""Chunk-autoregressive streaming generation.
+
+KVSampler works inside one packed span (rtcwm.model.packing.Window): chunk 0 runs together with the
+condition rows (fixed condition times, the span mask of rtcwm.model.masks); a finished chunk runs once more
+at model time 1 to add its keys/values (chunk 0 also the condition ones, kept per piece: text + Picture 1,
+each proxy chunk, audio), and a denoising step of chunk c >= 1 runs only its own rows against the cache --
+the static pieces, proxy chunks <= c (no future control) and the history. Equal to the full packed
+forward (tests/test_samplers.py).
+Step rules:
+    euler : the CWM flow sampler (velocity Euler steps), many steps
+    x0    : few-step (distilled students): predict the clean chunk, re-noise it to the next level
+
+generate_stream produces a whole clip chunk by chunk, with no restart. Stream chunk k is generated with
+its attention span, chunks j..k (j = span_start(k)), packed with that span's conditions and positions;
+the history chunks j..k-1 (already generated, clean) are re-run into the cache first. This is exactly the
+training situation of diffusion forcing with a clean prefix. kv_reuse=True instead keeps the history's keys
+and values across chunks (moved to the new positions), re-running only the conditions -- faster, approximate.
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+
+from cwm_h3_inference.constants import VIDEO_SHIFT
+
+from .layout import CLEAN_T, NUM_CHUNKS, SPAN_CHUNKS, chunk_row_slice, chunk_shape, span_start
+from .masks import df_mask
+from .objectives import make_window, shift
+from .packing import Window, run_blocks, video_output
+
+def schedule(steps: int, device, shift_value: float = VIDEO_SHIFT) -> torch.Tensor:
+    """steps + 1 noise levels from 1 to 0, flow-shifted."""
+    return shift(torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64, device=device), shift_value).float()
+
+
+class KVSampler:
+    """`train=True` builds the cache under no_grad (its tensors may then enter an autograd graph) and
+    checkpoints gradient steps; otherwise everything runs under inference_mode."""
+
+    def __init__(self, model, win: Window, train: bool = False):
+        self.model, self.win, self.train = model, win, train
+        self.ca_kv = None                      # per layer: condition pieces [head, proxy 0..5, audio], cached with chunk 0
+        self.chunks_kv = []                    # per finished chunk: [(k, v)] per layer; None once out of the window
+
+    def _no_grad(self):
+        return torch.no_grad() if self.train else torch.inference_mode()
+
+    def past(self, c: int):
+        if c == 0:
+            return None                        # chunk 0 runs with the condition rows themselves
+        hist = self.chunks_kv[max(0, c - (SPAN_CHUNKS - 1)):c]
+        assert all(h is not None for h in hist)
+        return [[self.ca_kv[li][0]] + self.ca_kv[li][1:2 + c] + [self.ca_kv[li][-1]] + [h[li] for h in hist]
+                for li in range(len(self.ca_kv))]
+
+    def step(self, c: int, latents, model_t: float, past=None, store=None, grad: bool = False):
+        """One forward of chunk c (latents [1,24,5,H,W]) at model time model_t -> its velocity."""
+        if grad:
+            # autocast caches the bf16 casts of the (fp32) LoRA weights for the whole autocast region; the
+            # casts made by earlier no-grad steps carry no autograd history and would cut the gradient
+            torch.clear_autocast_cache()
+        ctx = torch.enable_grad() if grad else self._no_grad()
+        with ctx:
+            win = self.win
+            hv = win.video_rows(self.model, latents)
+            t_rows = torch.full((hv.shape[1],), float(model_t))
+            tags = torch.zeros(hv.shape[1], dtype=torch.int64)
+            pos = win.target_pos[chunk_row_slice(c)]
+            if c == 0:
+                if getattr(self, "mask0", None) is None:
+                    self.mask0 = df_mask(win.head_rows, win.duv_rows, win.audio_rows, 1, win.dev,
+                                         static_sees_first=win.static_sees_first)
+                hidden, temb, inv = run_blocks(self.model, torch.cat([win.rows, hv], 1),
+                                               torch.cat([win.cond_times(), t_rows]), torch.cat([win.tags, tags]),
+                                               torch.cat([win.ca_pos, pos]), mask=self.mask0, train=grad,
+                                               kv=dict(past=None, store=store))
+                return video_output(self.model, hidden, temb, inv, win.rows.shape[1], latents.shape[2])
+            hidden, temb, inv = run_blocks(self.model, hv, t_rows, tags, pos, train=grad,
+                                           kv=dict(past=past if past is not None else self.past(c), store=store))
+            return video_output(self.model, hidden, temb, inv, 0, latents.shape[2])
+
+    def add_conditions(self, static=None):
+        """Cache the condition pieces alone: a span that does not start the stream, whose condition rows see
+        no video row (rtcwm.model.masks), so they come out exactly as in add_chunk(0, ...). static: the
+        static pieces [(head, audio)] per layer of an earlier such span -- the same in every one -- so only
+        the proxy rows are run, each proxy chunk seeing the static rows and the proxy chunks up to it."""
+        assert not self.win.static_sees_first
+        win, store = self.win, []
+        if static is None:
+            mask = df_mask(win.head_rows, win.duv_rows, win.audio_rows, 0, win.dev, static_sees_first=False)
+            with self._no_grad():
+                run_blocks(self.model, win.rows, win.cond_times(), win.tags, win.ca_pos, mask=mask,
+                           kv=dict(past=None, store=store))
+            self.ca_kv = [self._pieces(k, v) for k, v in store]
+            return
+        rows = slice(win.head_rows, win.head_rows + win.duv_rows)
+        dr = win.duv_rows // NUM_CHUNKS
+        with self._no_grad():
+            run_blocks(self.model, win.rows[:, rows], win.cond_times()[rows], win.tags[rows], win.ca_pos[rows],
+                       kv=dict(past=[[head, audio] for head, audio in static], store=store, causal_blocks=dr))
+        self.ca_kv = [[head] + [(k[:, :, a:a + dr].clone(), v[:, :, a:a + dr].clone()) for a in range(0, win.duv_rows, dr)]
+                      + [audio] for (head, audio), (k, v) in zip(static, store)]
+
+    def _pieces(self, k, v):
+        """Condition keys/values [B,H,rows,D] -> pieces [head, proxy 0..5, audio] (copies: views would keep
+        the whole run alive)."""
+        w = self.win
+        dr = w.duv_rows // NUM_CHUNKS
+        cuts = [0, w.head_rows] + [w.head_rows + (p + 1) * dr for p in range(NUM_CHUNKS)] + [w.rows.shape[1]]
+        return [(k[:, :, a:b].clone(), v[:, :, a:b].clone()) for a, b in zip(cuts[:-1], cuts[1:])]
+
+    def add_chunk(self, c: int, latents, past=None, model_t: float = CLEAN_T):
+        """Cache the keys/values of finished chunk c (run at model time model_t: 1 = clean)."""
+        store = []
+        self.step(c, latents.detach(), model_t, past=past, store=store)
+        assert len(store) == len(self.model.blocks)
+        if c == 0:                                            # split off the condition pieces
+            n = self.win.rows.shape[1]
+            self.ca_kv = []
+            for li in range(len(store)):                      # layer by layer
+                k, v = store[li]
+                self.ca_kv.append(self._pieces(k[:, :, :n], v[:, :, :n]))
+                store[li] = (k[:, :, n:].clone(), v[:, :, n:].clone())
+                del k, v
+        self.chunks_kv.append(store)
+        old = c + 1 - SPAN_CHUNKS                           # no later chunk looks at it any more
+        if old >= 0:
+            self.chunks_kv[old] = None
+
+    def generate_chunk(self, c: int, noise, sigmas, rule: str = "euler", gen=None, grad_step: int | None = None,
+                       cache: bool = True, exit_step: int | None = None, keep_last_kv: bool = False):
+        """Denoise chunk c from `noise` over `sigmas`. x0 rule only: with exit_step=k the chunk is the clean
+        prediction of step k (self forcing's random exit; later steps are skipped), with grad_step=k also
+        with its graph kept. cache: add the finished chunk to the cache (not needed when no later chunk of
+        this window follows). keep_last_kv: keep the keys/values the chunk's rows had in its last step (the
+        one its output comes from) in self.last_kv -- a history without the extra clean forward of
+        add_chunk (HISTORY_FROM_LAST_STEP)."""
+        past = self.past(c)
+        x, out = noise, None
+        stop = grad_step if grad_step is not None else exit_step
+        for i in range(len(sigmas) - 1):
+            s, s_next = float(sigmas[i]), float(sigmas[i + 1])
+            g = grad_step == i
+            last = i == stop or i == len(sigmas) - 2
+            store = [] if keep_last_kv and last else None
+            v = self.step(c, x, 1.0 - s, past=past, grad=g, store=store).float()
+            if store is not None:
+                self.last_kv = [(k.detach(), v_.detach()) for k, v_ in store]
+            if rule == "euler":
+                x = x + (s - s_next) * v
+                continue
+            x0 = x + s * v                                     # noisy = (1-s) x0 + s n, v = x0 - n
+            if i == stop or i == len(sigmas) - 2:
+                out = x0
+                break
+            with self._no_grad():
+                x = (1 - s_next) * x0 + s_next * torch.randn(x0.shape, generator=gen).to(x0.device)
+        out = x if rule == "euler" else out
+        if cache and c < NUM_CHUNKS - 1:
+            self.add_chunk(c, out, past=past)
+        return out
+
+
+ROTATE_IN_ATTENTION = True   # kv_reuse: cached keys keep their positions, moved by rotation inside attention
+HISTORY_FROM_LAST_STEP = True   # kv_reuse, sampling only: cache a chunk with the keys/values of its last denoising
+                                # step (noisy input) instead of one more clean forward -- one forward fewer per chunk,
+                                # and the noisy history damps the drift of a long rollout (a training rollout gets the
+                                # same history from rollout(history_noise=...))
+PROXY_ONLY_CONDITIONS = False   # kv_reuse: on a span move re-run only the proxy rows against the shared static
+                                # pieces (experimental, pending tests/test_conditions.py) instead of all conditions
+
+
+def _move(model, kv, delta: float, dev):
+    """[cached chunk, its pending time shift] after the span moved by delta. Normally the keys stay as they
+    were cached and only the shift accumulates (_moved hands it to attention), so every user of a chunk --
+    including the autograd graphs of a training rollout, which keep what they attended to until the
+    backward pass -- shares one copy. ROTATE_IN_ATTENTION = False re-rotates a copy of the keys at every
+    move (the earlier behaviour, kept for tests)."""
+    store, shift = kv
+    if not ROTATE_IN_ATTENTION:
+        return [_shift_time(model, store, delta, dev), 0.0]
+    return [store, shift + delta]
+
+
+def _moved(model, store, shift: float, dev, tables: dict):
+    """Per-layer pieces of a cached chunk at its current position: (k, v) or (k, v, rotation by shift)."""
+    if shift == 0:
+        return store
+    tab = tables.get(shift)
+    if tab is None:
+        tab = tables[shift] = model._rotation_table(torch.tensor([[[shift, 0.0, 0.0]]], device=dev), store[0][0].dtype)
+    return [(k, v, tab) for k, v in store]
+
+
+def _share_static(ca, static, sees_first: bool):
+    """Static condition pieces (caption + Picture 1, audio) of spans that do not start the stream see only
+    static rows at fixed positions, so they are the same in every such span: keep the first span's and
+    hand it to the later ones (the re-run's own copies are dropped), one copy instead of one per span."""
+    if sees_first:
+        return static
+    if static is None:
+        return [(layer[0], layer[-1]) for layer in ca]
+    for layer, (head, audio) in zip(ca, static):
+        layer[0], layer[-1] = head, audio
+    return static
+
+
+def _shift_time(model, store, delta: float, dev):
+    """Cached keys re-positioned by delta on the time axis: RoPE rotations compose, so rotating the keys
+    of position p by delta gives exactly the keys of position p + delta (values carry no position)."""
+    from .attention import rope
+    table = model._rotation_table(torch.tensor([[[delta, 0.0, 0.0]]], device=dev), store[0][0].dtype)
+    return [(rope(k.transpose(1, 2), table).transpose(1, 2).contiguous(), v) for k, v in store]
+
+
+@torch.inference_mode()
+def generate_stream(model, window_sample: Callable[[int], dict], n_chunks: int, steps: int, gen, dev,
+                    rule: str = "euler", shift_value: float = VIDEO_SHIFT, history_noise: float = 0.0,
+                    kv_reuse: bool = False, on_chunk=None) -> torch.Tensor:
+    """Stream chunks 0..n_chunks-1 -> [1,24,5*n_chunks,H,W]. window_sample(j) returns the inputs of the
+    span starting at stream chunk j (rtcwm.train.dataset.load_window); on_chunk(k, seconds) reports progress.
+    history_noise s > 0: the history chunks enter the cache noised to level s (model time 1 - s), as
+    diffusion forcing trains them -- damps the drift of a long rollout on its own output.
+    kv_reuse: keep each history chunk's keys/values from when it was cached and only move them to the
+    new span's positions (a rotation applied inside attention, _move); only the condition rows (whose proxy
+    changes with the span) are re-run, alone. Saves the re-run of the history, but its keys and values were
+    computed in the earlier spans -- an approximation of the exact default."""
+    out = _stream(model, window_sample, n_chunks, schedule(steps, dev, shift_value), gen, dev, rule, history_noise,
+                  kv_reuse, on_chunk=on_chunk)
+    return torch.cat(out, 2)
+
+
+def rollout(model, window_sample: Callable[[int], dict], n_chunks: int, sigmas, gen, dev, exit_steps,
+            grad_chunks=(), kv_reuse: bool = True, history_noise: float = 0.0, on_chunk=None) -> list[torch.Tensor]:
+    """Self-forcing rollout of a training student (x0 rule), the generate_stream procedure except that
+    chunk k stops at its step exit_steps[k] and is that step's clean prediction (and so the history of the
+    later chunks); the chunks in grad_chunks keep the graph of that step (the history they attend to enters
+    detached). Returns the chunks [1,24,5,H,W]. Every rank must pass the same exit_steps / grad_chunks:
+    they fix the sequence of block calls, and so of the FSDP gathers. history_noise s > 0: every history
+    chunk but the stream's first enters the cache noised to level s -- with s = sigmas[-2] the history a
+    HISTORY_FROM_LAST_STEP sampler caches (the last step's input), whatever step the chunk exited at."""
+    return _stream(model, window_sample, n_chunks, sigmas, gen, dev, "x0", history_noise, kv_reuse, on_chunk=on_chunk,
+                   train=True, exit_steps=exit_steps, grad_chunks=set(grad_chunks), clean_first=True)
+
+
+def _stream(model, window_sample, n_chunks, sig, gen, dev, rule, history_noise, kv_reuse, on_chunk=None,
+            train: bool = False, exit_steps=None, grad_chunks=frozenset(), clean_first: bool = False) -> list[torch.Tensor]:
+    import time
+    from musubi_tuner.minimax_h3.packing import _video_t_spans
+    chunk_span = sum(_video_t_spans(chunk_shape()[2]))       # time-axis length of one chunk
+    out, hist_lat, hist_kv, ca, prev_j = [], [], {}, None, None
+    static, tables = None, {}                                # see _share_static / _moved
+    for k in range(n_chunks):
+        t0 = time.time()
+        j = span_start(k)
+        with torch.no_grad():
+            win = make_window(model, window_sample(j), dev, gen)
+        kvs = KVSampler(model, win, train=train)
+        with kvs._no_grad():
+            if kv_reuse:
+                if prev_j is not None and j != prev_j:           # the span moved: drop its old first chunk,
+                    delta = -(j - prev_j) * chunk_span            # move the rest, re-run the conditions
+                    hist_kv = {i: _move(model, kv, delta, dev) for i, kv in hist_kv.items() if i >= j}
+                    ca = None
+                if k > j and ca is None:                         # (a moved span never starts the stream:
+                    kvs.add_conditions(static if PROXY_ONLY_CONDITIONS else None)  # its conditions see no video, and the
+                    ca = kvs.ca_kv                                # span's first chunk keeps its keys/values)
+                    static = _share_static(ca, static, win.static_sees_first)
+                kvs.ca_kv, kvs.chunks_kv = ca, [_moved(model, *hist_kv[i], dev, tables) for i in range(j, k)]
+            else:
+                for c in range(k - j):                             # the history of this span, re-run
+                    kvs.add_chunk(c, *hist_lat[j + c])
+        noise = torch.randn(chunk_shape(), generator=gen).to(dev)
+        e = exit_steps[k] if exit_steps is not None else None
+        last_kv = HISTORY_FROM_LAST_STEP and not train and kv_reuse and k > j and k < n_chunks - 1
+        x = kvs.generate_chunk(k - j, noise, sig, rule, gen, grad_step=e if k in grad_chunks else None, cache=False,
+                               exit_step=e, keep_last_kv=last_kv).float()
+        out.append(x)
+        with kvs._no_grad():
+            h, hn = x.detach(), 0.0 if clean_first and k == 0 else history_noise
+            if hn > 0:
+                h = (1 - hn) * h + hn * torch.randn(h.shape, generator=gen).to(dev)
+            hist_lat.append((h, 1.0 - hn))
+            if kv_reuse and k < n_chunks - 1:                      # cache the new chunk for the next spans
+                if k == j:                                        # (only k = 0: it runs with the conditions)
+                    kvs.add_chunk(0, h, model_t=1.0 - hn)
+                    ca, hist_kv[k] = kvs.ca_kv, [kvs.chunks_kv[0], 0.0]
+                    static = _share_static(ca, static, win.static_sees_first)
+                elif last_kv:
+                    hist_kv[k] = [kvs.last_kv, 0.0]
+                else:
+                    kvs.add_chunk(k - j, h, model_t=1.0 - hn)
+                    hist_kv[k] = [kvs.chunks_kv[-1], 0.0]
+        prev_j = j
+        del kvs, win
+        if on_chunk:
+            on_chunk(k, time.time() - t0)
+    return out
diff --git a/rtcwm/paths.py b/rtcwm/paths.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b35aadf828b76ca44ac4b62b7d88eaea81ff8a5
--- /dev/null
+++ b/rtcwm/paths.py
@@ -0,0 +1,51 @@
+"""Every filesystem location the project uses. Override with environment variables.
+
+    RTCWM_DATA    renders, pre-encoded cache, QC reports, runs, eval outputs (default: repo root)
+    CWM_RELEASE   the released Code World Model ("release/" of the CWM repo)
+    OCWM          open-code-world-model checkout (coarse proxy renderer, rig, follow camera)
+    QWEN_VL       Qwen3-VL-32B-Instruct snapshot used for captions
+"""
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+REPO = Path(__file__).resolve().parents[1]
+DATA = Path(os.environ.get("RTCWM_DATA", REPO))
+CWM_RELEASE = Path(os.environ.get("CWM_RELEASE", Path.home() / "projects/code-world-model/release"))
+OCWM = Path(os.environ.get("OCWM", Path.home() / "projects/open-code-world-model"))
+
+# --- data generation
+CARLA_ROOT = Path(os.environ.get("CARLA_ROOT", DATA / "CARLA_0.9.15"))
+DATASET = DATA / "dataset_v1"                      # one directory per rendered clip
+QC = DATA / "qc"                                   # per-clip QC reports (clipping/, audit/, coarse/)
+DATA_LISTS = REPO / "data_lists"                   # clips kept and windows excluded (in git)
+
+# --- training
+CACHE = DATA / "train_cache"                       # pre-encoded windows, one directory per clip
+RUNS = DATA / "runs"
+EVAL = DATA / "eval"
+
+# --- released models
+MODELS = CWM_RELEASE / "models"
+DIT = MODELS / "diffusion_models/minimax_h3_ref2va_bf16.safetensors"
+CWM_LORA = MODELS / "lora/minimax-h3-cwm.safetensors"
+VIDEO_VAE = MODELS / "vae/minimax_h3_video_vae_fp16.safetensors"
+AUDIO_VAE = MODELS / "vae/minimax_h3_audio_vae_fp32.safetensors"
+TEXT_ENCODER = MODELS / "text_encoders/qwen3vl_32b_minimax_h3_bf16.safetensors"
+PROCESSOR = MODELS / "processors/Qwen3-VL-32B-Instruct"
+CWM_EXAMPLES = CWM_RELEASE / "examples"
+
+
+def qwen_vl() -> Path:
+    if "QWEN_VL" in os.environ:
+        return Path(os.environ["QWEN_VL"])
+    snaps = Path.home() / ".cache/huggingface/hub/models--Qwen--Qwen3-VL-32B-Instruct/snapshots"
+    return next(snaps.iterdir())
+
+
+def use_ocwm() -> None:
+    """Make open-code-world-model's `compiler` package importable (it is not pip-installable)."""
+    if str(OCWM) not in sys.path:
+        sys.path.insert(0, str(OCWM))
diff --git a/rtcwm/train/__init__.py b/rtcwm/train/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d94fbd6823d142c6042f0deaab1b4ee6bdfc624
--- /dev/null
+++ b/rtcwm/train/__init__.py
@@ -0,0 +1 @@
+"""Training: pre-encoding rendered windows into model inputs, and causal LoRA training (`cwm` environment)."""
diff --git a/rtcwm/train/dataset.py b/rtcwm/train/dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a30f1b647f2454663b521b4ab68fd2e9d4c60df
--- /dev/null
+++ b/rtcwm/train/dataset.py
@@ -0,0 +1,117 @@
+"""Training samples: spans of 6 consecutive chunks cut from a clip's latent stream, each with a start frame.
+
+A clip's stream (rtcwm.train.preencode --stage stream) holds K chunks (17-frame VAE clips, K = 54 for the
+934-frame renders). The span starting at chunk j covers frames 17j .. 17j+101 and is packed with
+    x / DUV   : the stream's latents 5j .. 5j+29
+    Picture 1 : a start frame m <= j (rtcwm.data.anchors: the first frame of every 3rd chunk; native
+                1344x768 by default), with its static caption's Qwen states (that frame only, no timing)
+    audio     : silent
+With probability P_START the span starts at its start frame (j = m: generation from that image, no
+history; stream_start, so the static rows see chunk 0); otherwise the start frame is an earlier one
+(the stream has been running since m). The sample list is every span start of the kept clips
+(data_lists/clips.json: no Town02 / Town04) minus spans overlapping a caption window flagged by QC
+(data_lists/bad_windows.json). A fixed set of clips is held out.
+"""
+from __future__ import annotations
+
+import random
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+from cwm_h3_inference.constants import AUDIO_LATENT_SHAPE, STRIDE_FRAMES, VIDEO_LATENT_SHAPE, WINDOW_FRAMES
+
+from .. import paths
+from ..data.anchors import cache_paths, clip_anchors
+from ..data.qc import lists
+from ..model.layout import CHUNK_LATENTS, CLIP_FRAMES, TARGET_LATENTS, SPAN_CHUNKS
+from ..model.packing import audio_frames
+
+P_START = 0.2
+
+
+def stream_chunks(clip: str) -> int:
+    with safe_open(str(paths.CACHE / clip / "stream.safetensors"), "pt") as f:
+        return f.get_slice("duv").get_shape()[1] // CHUNK_LATENTS
+
+
+def start_frames(clip: str, res: str = "native") -> list[int]:
+    """Start chunks of a clip whose Picture 1 and caption states are encoded (at Picture 1 resolution res)."""
+    return [m for m in clip_anchors(clip)
+            if cache_paths(clip, m, res)["text"].exists() and cache_paths(clip, m, res)["latent"].exists()]
+
+
+def list_samples(val_clips: int, seed: int, res: str = "native"):
+    """(train, val, n_dropped) as lists of (clip, span start chunk)."""
+    keep, bad = lists.load()
+    samples, dropped = [], 0
+    for c in sorted(keep):
+        if not (paths.CACHE / c / "stream.safetensors").exists() or 0 not in start_frames(c, res):
+            continue
+        flagged = [(STRIDE_FRAMES * w, STRIDE_FRAMES * w + WINDOW_FRAMES) for w in bad.get(c, ())]
+        for j in range(stream_chunks(c) - SPAN_CHUNKS + 1):
+            a, b = CLIP_FRAMES * j, CLIP_FRAMES * (j + SPAN_CHUNKS)
+            if any(a < fb and fa < b for fa, fb in flagged):
+                dropped += 1
+                continue
+            samples.append((c, j))
+    names = sorted({c for c, _ in samples})
+    random.Random(seed).shuffle(names)
+    val = set(names[:val_clips])
+    return [s for s in samples if s[0] not in val], [s for s in samples if s[0] in val], dropped
+
+
+class StartFrames:
+    """Draws (span start j, start frame m) for a listed (clip, j): with probability P_START a span that
+    starts at one of the clip's start frames, otherwise the listed span with an earlier start frame."""
+
+    def __init__(self, samples, seed: int, res: str = "native"):
+        self.rng = random.Random(seed)
+        self.spans = {}
+        for c, j in samples:
+            self.spans.setdefault(c, set()).add(j)
+        self.frames = {c: start_frames(c, res) for c in self.spans}
+
+    def draw(self, clip: str, j: int) -> tuple[int, int]:
+        starts = [m for m in self.frames[clip] if m in self.spans[clip]]
+        if starts and self.rng.random() < P_START:
+            m = self.rng.choice(starts)
+            return m, m
+        earlier = [m for m in self.frames[clip] if m < j]
+        if earlier:
+            return j, self.rng.choice(earlier)
+        m = max(m for m in self.frames[clip] if m <= j)          # j = 0, or no start frame before it
+        return m, m
+
+
+def silence() -> torch.Tensor:
+    a = load_file(paths.CACHE / "silence_audio.safetensors")["audio"]
+    assert tuple(a.shape) == AUDIO_LATENT_SHAPE[1:]
+    return a
+
+
+def load_window(clip: str, j: int, silence_audio: torch.Tensor, start: int = 0, with_target: bool = True,
+                res: str = "native", n_latents: int = TARGET_LATENTS) -> dict:
+    """The span starting at stream chunk j, with start frame (chunk) start <= j: x [1,24,30,48,84]
+    (with_target), DUV ref [1,24,30,12,22], Picture 1, silent audio, Qwen states, stream_start.
+    n_latents = 37: the released window over the same frames (124 frames from chunk j; the stream's per-clip
+    latents are exactly its encoding, the VAE being causal inside a clip), with the released audio length."""
+    assert start <= j, (start, j)
+    d = paths.CACHE / clip
+    lat = slice(CHUNK_LATENTS * j, CHUNK_LATENTS * j + n_latents)
+    with safe_open(str(d / "stream.safetensors"), "pt") as f:
+        duv = f.get_slice("duv")[:, lat]
+        x = f.get_slice("target")[:, lat].unsqueeze(0) if with_target else None
+    assert duv.shape[1] == n_latents, (clip, j, duv.shape)
+    cp = cache_paths(clip, start, res)
+    t = load_file(cp["text"])
+    anchor = load_file(cp["latent"])["latents_ref_000_image"].unsqueeze(0)
+    ref = torch.nn.functional.pad(duv.unsqueeze(0), (0, 1, 0, 0), value=0.0)          # to the runtime width
+    n_audio = AUDIO_LATENT_SHAPE[-1] if n_latents == VIDEO_LATENT_SHAPE[2] else audio_frames(n_latents)
+    audio = silence_audio[..., :n_audio].unsqueeze(0).clone()
+    out = dict(audio=audio, anchor=anchor, ref=ref, hidden=t["hidden"].unsqueeze(0), tags=t["tags"].unsqueeze(0),
+               stream_start=j == start)
+    if with_target:
+        out["x"] = x
+    return out
diff --git a/rtcwm/train/distill.py b/rtcwm/train/distill.py
new file mode 100644
index 0000000000000000000000000000000000000000..dee2930e075693a5af22524845ed8fd147905610
--- /dev/null
+++ b/rtcwm/train/distill.py
@@ -0,0 +1,336 @@
+"""Few-step distillation of the streaming causal CWM by distribution matching (DMD) on the student's own
+long rollouts (self forcing), as in CausVid / Self Forcing / Self Forcing++.
+
+    torchrun ... -m rtcwm.train.distill --out DIR --student-lora DF_LORA [--steps 4] [--iters 200]
+
+Three LoRAs on one frozen, FSDP-sharded H3 (rtcwm.model.loading.LoRAStack):
+  student  the streaming DF model (rtcwm.train.train), sampled chunk by chunk with `--steps` x0-steps
+  teacher  the released CWM LoRA, frozen: the real score
+  critic   initialised from the released CWM LoRA, trained to denoise the student's samples: the fake score
+The teacher and the critic run the released model exactly as released: bidirectional over a released
+37-latent window (124 frames = 7 VAE clips + the first 2 latents of the 8th, a 5n+2 window), with that
+window's whole proxy. The VAE is causal inside a clip, so 37 consecutive latents of a chunk stream are
+exactly the released encoding of those frames (ops/dbg/tail_latents.py) -- no layout mismatch.
+
+Each iteration (rollout length R shared by all ranks, drawn from --min-chunks..--max-chunks):
+  1. every rank draws a clip and a start frame m and streams R chunks from it (rtcwm.model.sampler.rollout,
+     the inference procedure, with the KV cache carried across chunks). As in Self Forcing (mingjia's
+     port, wt-sf-mingjia/sf/self_forcing_model.py), every chunk stops at its own random x0-step (shared by
+     all ranks) and that step's clean prediction is the chunk -- and the history of the next ones. The
+     last --grad-chunks chunks of the teacher window keep the graph of their stopping step, their history
+     entering detached (memory: a graph keeps what it attended to -- its span's proxy keys/values and the
+     history chunks that have since left the window -- alive until the backward pass; four fit).
+     Earlier chunks keep no graph: the window slides, and a graph would keep its own shifted copy of the
+     history's keys/values alive until the backward pass.
+     A history chunk enters the cache as the sampler caches it (rtcwm.model.sampler.HISTORY_FROM_LAST_STEP):
+     noised to the last sampling step's level, the first chunk clean (--rollout-history last-step), or clean.
+  2. teacher window = chunks R-8 .. R-2 and the first 2 latents of chunk R-1 (Picture 1 = the window's own
+     first frame from the rollout, as the released CWM is used; the start frame's caption; the window's
+     proxy; silent audio), noised to one random level:
+        loss = 0.5 || x - sg(x - (x0_fake - x0_real) / mean|x - x0_real|) ||^2  over the window  (DMD gradient)
+  3. critic: --critic-per-gen flow-matching updates (generator : critic = 1 : 7 by default), each on a
+     random teacher window of the same detached rollout (R-7 of them) with fresh noise -- the rollout, the
+     expensive part, is shared.
+With --teacher-kind causal the teacher is a streaming DF checkpoint (--teacher-lora, frozen) and the critic
+starts from it: the teacher window is then the student's own layout -- one 6-chunk span of the rollout
+(the last one for the generator, a random one for the critic) with the stream's start frame and caption and
+the span's causal proxy (rule E, rtcwm.model.masks) -- every chunk of it noised to the same level.
+Long rollouts put the student in the drifted states it reaches at test time, and the teacher's score pulls
+them back -- training on its own history removes the exposure bias of diffusion forcing.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import time
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from cwm_h3_inference.constants import VIDEO_LATENT_SHAPE, VIDEO_SHIFT
+
+from .. import paths
+from ..model import loading
+from ..model.layout import CHUNK_LATENTS, SPAN_CHUNKS
+from ..model.objectives import df_forward, make_window, shift
+from ..model.packing import Window, run_blocks, video_output
+from ..model.sampler import rollout, schedule
+from .dataset import list_samples, load_window, silence, start_frames, stream_chunks
+from .train import node_group
+
+STUDENT_SHIFT = 5.0                    # few-step schedule: sigma 1, .94, .83, .63, 0 for 4 steps
+TEACHER_LATENTS = VIDEO_LATENT_SHAPE[2]                       # 37 = 7 chunks + 2
+TEACHER_CHUNKS = TEACHER_LATENTS // CHUNK_LATENTS + 1         # chunks a teacher window touches: 8
+
+
+def parse():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--student-lora", type=Path, required=True, help="streaming DF checkpoint to distil")
+    ap.add_argument("--teacher-lora", type=Path, default=paths.CWM_LORA)
+    ap.add_argument("--critic-lora", type=Path, default=paths.CWM_LORA)
+    ap.add_argument("--steps", type=int, default=4, help="student sampling steps")
+    ap.add_argument("--iters", type=int, default=200)
+    ap.add_argument("--min-chunks", type=int, default=TEACHER_CHUNKS, help="shortest rollout (chunks)")
+    ap.add_argument("--max-chunks", type=int, default=14, help="longest rollout (chunks; 14 = 10 s)")
+    ap.add_argument("--kv-reuse", choices=["on", "off"], default="on", help="rollout: carry the history's keys/values")
+    ap.add_argument("--rollout-history", choices=["last-step", "clean"], default="last-step",
+                    help="rollout history: last-step = noised to the last sampling step's level, the history the sampler "
+                         "caches by default (HISTORY_FROM_LAST_STEP); clean = one more clean forward")
+    ap.add_argument("--gen-lr", type=float, default=2e-6)
+    ap.add_argument("--critic-lr", type=float, default=4e-6)
+    ap.add_argument("--critic-per-gen", type=int, default=7, help="critic updates per generator update")
+    ap.add_argument("--grad-chunks", type=int, default=4, help="last teacher-window chunks that keep their graph")
+    ap.add_argument("--teacher-kind", choices=["released", "causal"], default="released",
+                    help="released: the released CWM on 37-latent bidirectional windows; causal: a streaming DF "
+                         "checkpoint (--teacher-lora, also the critic's start) on the student's 6-chunk spans")
+    ap.add_argument("--teacher-picture1", choices=["window", "start"], default="window",
+                    help="teacher / critic Picture 1: the window's first frame (rollout) or the stream's start frame")
+    ap.add_argument("--warmup", type=int, default=10)
+    ap.add_argument("--save-every", type=int, default=20)
+    ap.add_argument("--val-clips", type=int, default=12)
+    ap.add_argument("--picture1", choices=["native", "2048"], default="native")
+    ap.add_argument("--seed", type=int, default=0)
+    return ap.parse_args()
+
+
+def park(opts, device):
+    """Move the optimizers' moments to `device`: the CPU while the rollout's graphs fill the GPU (about
+    10 GB for student + critic), back before a step."""
+    for opt in opts:
+        for state in opt.state.values():
+            for key, val in state.items():
+                if torch.is_tensor(val) and val.dim() > 0:
+                    state[key] = val.to(device, non_blocking=True)
+
+
+class Reducer:
+    """Asynchronous gradient all-reduce for replicated LoRA weights, overlapped with backward."""
+
+    def __init__(self, params):
+        self.handles = []
+        for p in params:
+            p.register_post_accumulate_grad_hook(
+                lambda q: self.handles.append(dist.all_reduce(q.grad, op=dist.ReduceOp.AVG, async_op=True)))
+
+    def wait(self):
+        for h in self.handles:
+            h.wait()
+        self.handles.clear()
+
+
+class Starts:
+    """(clip, start frame m) pairs whose rollout of n chunks stays inside the training spans (no held-out
+    clip, no QC-flagged span) and whose teacher window fits the stream."""
+
+    def __init__(self, samples, res: str):
+        spans = {}
+        for c, j in samples:
+            spans.setdefault(c, set()).add(j)
+        self.spans, self.frames = spans, {c: start_frames(c, res) for c in spans}
+        self.n_stream = {c: stream_chunks(c) for c in spans}
+
+    def draw(self, rng: random.Random, n: int):
+        while True:
+            clip = rng.choice(sorted(self.spans))
+            ok = [m for m in self.frames[clip] if m + n <= self.n_stream[clip]
+                  and all(m + j in self.spans[clip] for j in range(0, n - SPAN_CHUNKS + 1))]
+            if ok:
+                return clip, rng.choice(ok)
+
+
+def released_velocity(model, win: Window, latents, model_t: float, train: bool):
+    """The released (bidirectional) H3 on a released window: latents [1,24,37,H,W] at model time model_t."""
+    hv = win.video_rows(model, latents)
+    t_rows = torch.cat([win.cond_times(), torch.full((hv.shape[1],), float(model_t))])
+    tags = torch.cat([win.tags, torch.zeros(hv.shape[1], dtype=torch.int64)])
+    hidden, temb, inv = run_blocks(model, torch.cat([win.rows, hv], 1), t_rows, tags,
+                                   torch.cat([win.ca_pos, win.target_pos]), train=train)
+    return video_output(model, hidden, temb, inv, win.rows.shape[1], latents.shape[2])
+
+
+def causal_velocity(model, win: Window, latents, model_t: float, train: bool):
+    """A streaming DF model on one span (rtcwm.model.objectives.df_forward), every chunk at model time model_t."""
+    return df_forward(model, win, latents, [float(model_t)] * SPAN_CHUNKS, {}, train=train)
+
+
+def noised(x, gen):
+    """x at a random teacher-schedule level -> (x_t, n, sigma)."""
+    sigma = float(shift(torch.rand(1, generator=gen).clamp(0.02, 0.98), VIDEO_SHIFT))
+    n = torch.randn(x.shape, generator=gen).to(x.device)
+    return (1 - sigma) * x + sigma * n, n, sigma
+
+
+def dmd_loss(stack, model, win, x, gen, velocity=released_velocity):
+    """DMD generator loss on teacher window x (the chunks that kept their graph carry the gradient)."""
+    xt, _, sigma = noised(x.detach(), gen)
+    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+        stack.use("teacher")
+        x0_real = xt + sigma * velocity(model, win, xt, 1.0 - sigma, train=False).float()
+        stack.use("critic")
+        x0_fake = xt + sigma * velocity(model, win, xt, 1.0 - sigma, train=False).float()
+    p = x.detach() - x0_real
+    grad = (x0_fake - x0_real) / p.abs().mean().clamp_min(1e-6)
+    stack.use("student")                                         # the graph of x is the student's
+    return 0.5 * torch.nn.functional.mse_loss(x.float(), (x.detach() - grad).float()), float(grad.abs().mean())
+
+
+def critic_loss(stack, model, win, x, gen, velocity=released_velocity):
+    """Flow-matching loss of the critic on a detached window."""
+    xt, n, sigma = noised(x.detach(), gen)
+    stack.use("critic")
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        v = velocity(model, win, xt, 1.0 - sigma, train=True)
+    return torch.nn.functional.mse_loss(v.float(), (x.detach() - n).float())
+
+
+def resume_point(a):
+    """(iterations done, student LoRA, critic LoRA): the latest student/critic pair saved in --out (a requeued
+    or resubmitted run continues from it -- lowprio jobs are preempted and requeued), else the initial ones."""
+    done = sorted(int(p.stem.split("iter")[1]) for p in a.out.glob("student_iter*.safetensors")
+                  if (a.out / p.name.replace("student", "critic")).exists())
+    if not done:
+        return 0, a.student_lora, a.critic_lora
+    n = done[-1]
+    return n, a.out / f"student_iter{n:05d}.safetensors", a.out / f"critic_iter{n:05d}.safetensors"
+
+
+def main():
+    a = parse()
+    causal = a.teacher_kind == "causal"
+    velocity = causal_velocity if causal else released_velocity
+    w_chunks = SPAN_CHUNKS if causal else TEACHER_CHUNKS            # rollout chunks one teacher window covers
+    assert a.grad_chunks <= w_chunks, (a.grad_chunks, w_chunks)
+    assert TEACHER_CHUNKS <= a.min_chunks <= a.max_chunks, (a.min_chunks, a.max_chunks)
+    dist.init_process_group("nccl")
+    rank, world = dist.get_rank(), dist.get_world_size()
+    local = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(local)
+    dev = torch.device("cuda", local)
+    log = (lambda *m: print(f"[{time.strftime('%H:%M:%S')}]", *m, flush=True)) if rank == 0 else (lambda *m: None)
+    group = node_group(rank, world)
+    # NCCL sets up a communicator's buffers (outside the PyTorch allocator) at its first collective. The
+    # first cross-node one would be the generator's gradient all-reduce, with the GPU at its fullest: make
+    # it now, and keep a few GB out of PyTorch's reach for NCCL.
+    torch.cuda.set_per_process_memory_fraction(0.96, dev)
+    dist.all_reduce(torch.zeros(1 << 24, device=dev))
+
+    train, val, _ = list_samples(a.val_clips, a.seed, a.picture1)
+    starts = Starts(train, a.picture1)
+    log(f"distill world={world} train_spans={len(train)} clips={len(starts.spans)} val={len(val)} iters={a.iters} "
+        f"student_steps={a.steps} rollout={a.min_chunks}..{a.max_chunks} chunks")
+    if rank == 0:
+        a.out.mkdir(parents=True, exist_ok=True)
+        (a.out / "val_samples.json").write_text(json.dumps(val))
+        (a.out / "args.json").write_text(json.dumps({k: str(v) for k, v in vars(a).items()}, indent=1))
+
+    start, student_lora, critic_lora = resume_point(a)
+    if start:
+        log(f"resuming after iter {start}: {student_lora.name}, {critic_lora.name} (optimizer moments start afresh)")
+    tr = loading.load_base()
+    stack = loading.LoRAStack(tr, dict(student=(student_lora, True), critic=(critic_lora, True),
+                                       teacher=(a.teacher_lora, False)))
+    loading.shard_blocks(tr, group, dev)
+    stack.to(dev)
+    tr.train()
+    gen_params, critic_params = stack.params("student"), stack.params("critic")
+    reducer = Reducer(gen_params + critic_params)
+    opt_g = torch.optim.AdamW(gen_params, lr=a.gen_lr, betas=(0.0, 0.999), weight_decay=0.01)
+    opt_c = torch.optim.AdamW(critic_params, lr=a.critic_lr, betas=(0.0, 0.999), weight_decay=0.01)
+    warm = lambda i: min(1.0, (i + 1) / a.warmup)                                   # constant after warmup
+    sched_g, sched_c = (torch.optim.lr_scheduler.LambdaLR(o, warm) for o in (opt_g, opt_c))
+    sigmas = schedule(a.steps, dev, STUDENT_SHIFT)
+    history_noise = float(sigmas[-2]) if a.rollout_history == "last-step" else 0.0
+    audio = silence()
+    gen = torch.Generator().manual_seed(a.seed * 1000 + rank)
+    rng = random.Random(a.seed * 1000 + rank)
+    log(f"student/critic trainable={sum(p.numel() for p in gen_params) / 1e6:.0f}M each, "
+        f"schedule={[round(float(s), 3) for s in sigmas]}")
+
+    def step(opt, sched, params):
+        reducer.wait()
+        gn = torch.nn.utils.clip_grad_norm_(params, 1.0)
+        opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
+        return float(gn)
+
+    t0 = time.time()
+    for it in range(start, a.iters):
+        # the rollout length and the gradient step fix the sequence of FSDP gathers: equal on every rank
+        shared = random.Random(a.seed * 100003 + it)
+        n = shared.randint(a.min_chunks, a.max_chunks)
+        exit_steps = [shared.randrange(a.steps) for _ in range(n)]
+        grad_chunks = list(range(n - a.grad_chunks, n))
+        clip, m = starts.draw(rng, n)
+
+        t_it, chunk_times = time.time(), []
+        park((opt_g, opt_c), "cpu")
+        stack.use("student")
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            chunks = rollout(tr, lambda j: load_window(clip, m + j, audio, start=m, with_target=False, res=a.picture1),
+                             n, sigmas, gen, dev, exit_steps, grad_chunks, kv_reuse=a.kv_reuse == "on", history_noise=history_noise,
+                             on_chunk=(lambda k, s: chunk_times.append(s)) if it < 3 else None)
+        t_roll = time.time() - t_it
+        if chunk_times:
+            log(f"iter {it + 1} rollout seconds per chunk: {[round(s, 1) for s in chunk_times]}")
+
+        def teacher_window(w):
+            """Rollout chunks w .. w+7 as a released window (37 latents) with its conditions. With
+            --teacher-picture1 window the window's Picture 1 is its own first frame -- the first latent of the
+            rollout's chunk w, which the per-clip VAE makes the frame's own encoding to within a few percent --
+            as the released CWM is used (Picture 1 = the window's first frame): the scores then ask for a
+            continuation of the scene the rollout is in, not of the start frame seen long ago.
+            Causal teacher: rollout chunks w .. w+5 as the student's span m+w (start frame, its caption, causal
+            proxy)."""
+            if causal:
+                sample = load_window(clip, m + w, audio, start=m, with_target=False, res=a.picture1)
+                with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+                    win = make_window(tr, sample, dev, gen)
+                return win, torch.cat(chunks[w:w + SPAN_CHUNKS], 2)
+            sample = load_window(clip, m + w, audio, start=m, with_target=False, res=a.picture1,
+                                 n_latents=TEACHER_LATENTS)
+            if a.teacher_picture1 == "window":
+                sample["anchor"] = chunks[w][:, :, :1].detach().cpu()
+            with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+                win = make_window(tr, sample, dev, gen)
+            return win, torch.cat(chunks[w:w + TEACHER_CHUNKS], 2)[:, :, :TEACHER_LATENTS]
+
+        win, x = teacher_window(n - w_chunks)
+        lg, gmag = dmd_loss(stack, tr, win, x, gen, velocity)
+        lg.backward()
+        park((opt_g, opt_c), dev)
+        gn_g = step(opt_g, sched_g, gen_params)
+        del win, x
+        chunks = [c.detach() for c in chunks]
+        lcs = []
+        for _ in range(a.critic_per_gen):
+            win, x = teacher_window(rng.randrange(n - w_chunks + 1))
+            lc = critic_loss(stack, tr, win, x, gen, velocity)
+            lc.backward()
+            gn_c = step(opt_c, sched_c, critic_params)
+            lcs.append(float(lc))
+        stats = torch.tensor([float(lg), sum(lcs) / len(lcs), gmag], device=dev)
+        dist.all_reduce(stats, op=dist.ReduceOp.AVG)
+        if rank == 0:
+            el = time.time() - t0
+            log(f"iter {it + 1}/{a.iters} R={n} grad_chunks={len(grad_chunks)} dmd={stats[0]:.4f} critic={stats[1]:.4f} "
+                f"|grad|={stats[2]:.3f} gn_g={gn_g:.2e} gn_c={gn_c:.2e} mem={torch.cuda.max_memory_allocated() / 2**30:.0f}G "
+                f"rollout {t_roll:.0f}s, iter {time.time() - t_it:.0f}s, eta={(a.iters - it - 1) * el / (it + 1 - start) / 3600:.2f}h")
+            with open(a.out / "train_log.jsonl", "a") as fh:
+                fh.write(json.dumps(dict(iter=it + 1, chunks=n, exit_steps=exit_steps, grad_chunks=grad_chunks, dmd=float(stats[0]),
+                                         critic=float(stats[1]), grad=float(stats[2]), gn_g=gn_g, gn_c=gn_c,
+                                         rollout_sec=t_roll, sec=el)) + "\n")
+            if (it + 1) % a.save_every == 0 or it + 1 == a.iters:
+                for name in ("student", "critic"):
+                    p = a.out / f"{name}_iter{it + 1:05d}.safetensors"
+                    stack.nets[name].save_weights(str(p), torch.bfloat16, {"base": "cwm", "iter": str(it + 1),
+                                                                         "objective": f"dmd-{a.steps}step"})
+                log("saved", a.out / f"student_iter{it + 1:05d}.safetensors")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/train/preencode.py b/rtcwm/train/preencode.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5e2720e05dea7191ccf5611c06367221644d7ba
--- /dev/null
+++ b/rtcwm/train/preencode.py
@@ -0,0 +1,317 @@
+"""Pre-encode rendered clips into training windows, exactly as the released CWM inference builds them.
+
+Per clip, in paths.CACHE/<clip>/:
+    anchor.safetensors        latents_ref_000_image      Picture 1 (frame 0, short edge 2048)
+    first_frame.safetensors   first_frame [24,1,48,84]   frame 0 at 1344x768 (window-0 keyframe)
+    wNN.visual.safetensors    target [24,37,48,84]       VAE latent of the window's 124 GT frames
+                              duv [24,37,12,21]          coarse-proxy reference latent
+                              prefix [24,10,48,84]       (NN>0) GT frames 0..33 re-encoded (Retake34)
+    wNN.text.safetensors      hidden [L,5120] bf16, tags [L]   Qwen conditioning (w0 / wn system prompt)
+    stream.safetensors        target [24,5K,48,84]       the whole clip as one latent stream: its first K
+                                                         17-frame VAE clips (918 of 934 frames), 5 latents each
+                              duv [24,5K,12,21]          the coarse proxy, encoded the same way
+    anchor_cMM.png / .safetensors, static_cMM.text.safetensors
+                              per start frame m (rtcwm.data.anchors): its Picture 1 latent and the Qwen states
+                              of its static caption (that frame and its proxy only, w0 system prompt)
+paths.CACHE/silence_audio.safetensors   audio [32,2,207]
+
+    python -m rtcwm.train.preencode --stage visual|text|stream|static_text [--shard] [--reverse] [CLIP_DIR ...]
+Without CLIP_DIRs every clip of the dataset with a complete prompt set is taken; --shard splits them
+over SLURM tasks (task i takes clips i, i+n, ...); --reverse walks from the end (a helper job meeting
+the main one in the middle). Existing files are skipped; writes are atomic per process.
+"""
+from __future__ import annotations
+
+import argparse
+import gc
+import os
+import time
+from pathlib import Path
+
+import av
+import numpy as np
+import torch
+from safetensors.torch import save_file
+
+from cwm_h3_inference.cache import _load_anchor
+from cwm_h3_inference.config import _canonical_caption
+from cwm_h3_inference.constants import (AUDIO_LATENT_SHAPE, FPS, OVERLAP_FRAMES, QWEN_OFFSETS, REFERENCE_LATENT_SHAPE,
+                                        STRIDE_FRAMES, TARGET_HEIGHT, TARGET_WIDTH, VIDEO_LATENT_SHAPE,
+                                        VIDEO_PREFIX_LATENTS, WINDOW_AUDIO_SAMPLES, WINDOW_FRAMES)
+from cwm_h3_inference.duv import load_duv_window
+from cwm_h3_inference.presentation import build_ref2va_presentation
+
+from .. import paths
+from ..data.clip import has_prompts, n_windows, qc_pass, read_meta
+
+
+def windows_of(clip: Path) -> int:
+    return min(n_windows(read_meta(clip)["frames"]), int(os.environ.get("MAX_WINDOWS", 99)))
+
+
+def read_frames(path: Path) -> np.ndarray:
+    with av.open(str(path)) as c:
+        return np.stack([f.to_ndarray(format="rgb24") for f in c.decode(video=0)])
+
+
+def atomic_save(path: Path, tensors: dict, meta: dict | None = None):
+    tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp")     # unique: a helper job may write the same file
+    save_file({k: v.contiguous() for k, v in tensors.items()}, str(tmp), metadata=meta or {})
+    tmp.replace(path)
+
+
+def stage_visual(clips):
+    from musubi_tuner.minimax_h3.audio_vae import load_audio_vae
+    from musubi_tuner.minimax_h3.video_vae import encode_video_condition, load_video_vae
+    from musubi_tuner.minimax_h3_generate_video import _prepare_pixels
+    from PIL import Image
+
+    dev = torch.device("cuda:0")
+    vae = load_video_vae(paths.VIDEO_VAE, device=dev, dtype=torch.float16, disable_mmap=False).eval().requires_grad_(False)
+    sil = paths.CACHE / "silence_audio.safetensors"
+    if not sil.exists():
+        avae = load_audio_vae(paths.AUDIO_VAE, device=dev, dtype=torch.float32, disable_mmap=False).eval()
+        with torch.inference_mode():
+            a = avae.encode(torch.zeros(1, 2, WINDOW_AUDIO_SAMPLES, device=dev)).float().cpu()[0]
+        assert tuple(a.shape) == AUDIO_LATENT_SHAPE[1:], a.shape
+        paths.CACHE.mkdir(parents=True, exist_ok=True)
+        atomic_save(sil, {"audio": a})
+        del avae
+
+    def enc(frames_uint8: np.ndarray) -> torch.Tensor:
+        px = _prepare_pixels(torch.from_numpy(np.ascontiguousarray(frames_uint8))).to(dev, torch.float16)
+        with torch.inference_mode():
+            return vae.encode(px).float().cpu()[0]
+
+    for clip in clips:
+        out = paths.CACHE / clip.name
+        out.mkdir(parents=True, exist_ok=True)
+        W = windows_of(clip)
+        if all((out / f"w{w:02d}.visual.safetensors").exists() for w in range(W)) and (out / "anchor.safetensors").exists():
+            continue
+        frames = read_frames(clip / "rgb.mp4")
+        assert frames.shape[1:3] == (TARGET_HEIGHT, TARGET_WIDTH), frames.shape
+        anchor_png = out / "first_frame.png"
+        Image.fromarray(frames[0]).save(anchor_png)
+        if not (out / "anchor.safetensors").exists():
+            big = _load_anchor(anchor_png)
+            px = _prepare_pixels(torch.from_numpy(big).unsqueeze(0)).to(dev, torch.float16)
+            with torch.inference_mode():
+                lat = encode_video_condition(vae, px)[0].cpu().float()
+            atomic_save(out / "anchor.safetensors", {"latents_ref_000_image": lat})
+            ff = enc(frames[:1])
+            assert ff.shape[1] == 1
+            atomic_save(out / "first_frame.safetensors", {"first_frame": ff})
+        for w in range(W):
+            p = out / f"w{w:02d}.visual.safetensors"
+            if p.exists():
+                continue
+            s = STRIDE_FRAMES * w
+            target = enc(frames[s:s + WINDOW_FRAMES])
+            assert tuple(target.shape) == VIDEO_LATENT_SHAPE[1:], target.shape
+            duv = load_duv_window(clip / "conditions_coarse", s, for_qwen=False)
+            dpx = _prepare_pixels(torch.from_numpy(duv.copy())).to(dev, torch.float16)
+            with torch.inference_mode():
+                ref = encode_video_condition(vae, dpx)[0].cpu().float()
+            assert tuple(ref.shape) == REFERENCE_LATENT_SHAPE, ref.shape
+            t = {"target": target, "duv": ref}
+            if w > 0:
+                tail = frames[s:s + OVERLAP_FRAMES]
+                padded = np.concatenate([tail, np.repeat(tail[-1:], WINDOW_FRAMES - OVERLAP_FRAMES, 0)], 0)
+                t["prefix"] = enc(padded)[:, :VIDEO_PREFIX_LATENTS]
+            atomic_save(p, t)
+            gc.collect(); torch.cuda.empty_cache()
+        print("visual", clip.name, W, "windows", flush=True)
+
+
+def stage_text(clips):
+    from musubi_tuner.minimax_h3.media import H3AudioSource, H3Record, H3Reference
+    from musubi_tuner.minimax_h3.text_encoder import (H3TextVisual, build_presentation, encode_h3_presentation,
+                                                       load_h3_processor, load_h3_text_encoder)
+    dev = torch.device("cuda:0")
+    proc = load_h3_processor(str(paths.PROCESSOR))
+    model = load_h3_text_encoder(str(paths.TEXT_ENCODER), processor_path=str(paths.PROCESSOR), device=dev,
+                                 dtype=torch.bfloat16).eval()
+    for clip in clips:
+        out = paths.CACHE / clip.name
+        anchor_png = out / "first_frame.png"
+        anchor = torch.from_numpy(np.ascontiguousarray(_load_anchor(anchor_png))).unsqueeze(0)
+        cond = clip / "conditions_coarse"
+        for w in range(windows_of(clip)):
+            p = out / f"w{w:02d}.text.safetensors"
+            if p.exists():
+                continue
+            caption = _canonical_caption((clip / "prompts" / f"{w:03d}.txt").read_text(), f"{clip.name}/w{w}")
+            s = STRIDE_FRAMES * w
+            duv = torch.from_numpy(np.ascontiguousarray(load_duv_window(cond, s, for_qwen=True)))
+            image_key, video_key = anchor_png, cond / f"{s:06d}.semantic_id.png"
+            placeholder = cond / ".h3-presentation-placeholder"
+            record = H3Record(video_path=placeholder, caption=caption,
+                              target_audio=H3AudioSource(path=placeholder, embedded=False),
+                              references=(H3Reference(type="image", path=image_key),
+                                          H3Reference(type="video", path=video_key, duration_seconds=WINDOW_FRAMES / FPS)),
+                              jsonl_line=0)
+            plain = build_presentation(record, "ref2va", {image_key: H3TextVisual(anchor),
+                                                          video_key: H3TextVisual(duv, tuple(o / FPS for o in QWEN_OFFSETS))})
+            pres = build_ref2va_presentation(proc, plain, w, caption)        # w0 / wn system prompt by index
+            with torch.inference_mode():
+                hidden, tags = encode_h3_presentation(proc, model, pres)
+            atomic_save(p, {"hidden": hidden.to("cpu", torch.bfloat16), "tags": tags.to("cpu", torch.int64)},
+                        {"system_prompt": "w0" if w == 0 else "wn"})
+            gc.collect(); torch.cuda.empty_cache()
+        print("text", clip.name, flush=True)
+
+
+CLIP_FRAMES = 17                       # the video VAE encodes independent 17-frame clips into 5 latents each
+
+
+def stage_stream(clips):
+    """The whole clip as a continuous latent stream (streaming generation). The VAE encodes 17-frame clips
+    independently, so the stream is the concatenation of per-clip latents; one dummy clip is appended
+    because the encoder drops the last 3 latents of whatever it is given."""
+    from musubi_tuner.minimax_h3.video_vae import encode_video_condition, load_video_vae
+    from musubi_tuner.minimax_h3_generate_video import _prepare_pixels
+    from cwm_h3_inference.duv import load_duv_frame
+
+    dev = torch.device("cuda:0")
+    vae = load_video_vae(paths.VIDEO_VAE, device=dev, dtype=torch.float16, disable_mmap=False).eval().requires_grad_(False)
+
+    def padded(frames: np.ndarray) -> torch.Tensor:
+        frames = np.concatenate([frames, np.repeat(frames[-1:], CLIP_FRAMES, 0)], 0)
+        return _prepare_pixels(torch.from_numpy(np.ascontiguousarray(frames))).to(dev, torch.float16)
+
+    for clip in clips:
+        p = paths.CACHE / clip.name / "stream.safetensors"
+        if p.exists():
+            continue
+        frames = read_frames(clip / "rgb.mp4")
+        k = len(frames) // CLIP_FRAMES
+        n = k * CLIP_FRAMES
+        duv = np.stack([load_duv_frame(clip / "conditions_coarse", i, for_qwen=False) for i in range(n)])
+        with torch.inference_mode():
+            target = vae.encode(padded(frames[:n])).float().cpu()[0, :, :5 * k]
+            ref = encode_video_condition(vae, padded(duv))[0].float().cpu()[:, :5 * k]
+        assert target.shape[1] == ref.shape[1] == 5 * k, (target.shape, ref.shape)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        atomic_save(p, {"target": target, "duv": ref}, {"clip_frames": str(CLIP_FRAMES), "vae_clips": str(k)})
+        gc.collect(); torch.cuda.empty_cache()
+        print("stream", clip.name, k, "VAE clips", flush=True)
+
+
+def stage_static_text(clips, follow_job: str = "", res: str = "native"):
+    """Per start frame (rtcwm.data.anchors): the Picture 1 latent of that frame (short edge 2048, as the
+    released CWM encodes its anchor) and the Qwen states of its static caption. The text must not carry
+    anything from later frames: Qwen sees Picture 1 and, as the Ref2VA video reference, the proxy of the
+    start frame repeated (not a window's proxy video), with the first-window system prompt."""
+    from musubi_tuner.minimax_h3.text_encoder import load_h3_processor, load_h3_text_encoder
+    from musubi_tuner.minimax_h3.video_vae import load_video_vae
+    dev = torch.device("cuda:0")
+    vae = load_video_vae(paths.VIDEO_VAE, device=dev, dtype=torch.float16, disable_mmap=False).eval().requires_grad_(False)
+    proc = load_h3_processor(str(paths.PROCESSOR))
+    model = load_h3_text_encoder(str(paths.TEXT_ENCODER), processor_path=str(paths.PROCESSOR), device=dev,
+                                 dtype=torch.bfloat16).eval()
+    while True:
+        running = bool(follow_job) and _job_running(follow_job)      # captions still being written?
+        n = _static_text_pass(clips, vae, proc, model, dev, res)
+        if not running:
+            break
+        if not n:
+            time.sleep(60)
+
+
+def _static_text_pass(clips, vae, proc, model, dev, res: str = "native") -> int:
+    from PIL import Image
+    from musubi_tuner.minimax_h3.video_vae import encode_video_condition
+    from musubi_tuner.minimax_h3_generate_video import _prepare_pixels
+    from ..data.anchors import anchor_frame, cache_paths, caption_path, clip_anchors
+    done = 0
+    for clip in clips:
+        todo = [m for m in clip_anchors(clip) if caption_path(clip, m).exists() and not cache_paths(clip.name, m, res)["text"].exists()]
+        if not todo:
+            continue
+        frames = read_frames(clip / "rgb.mp4")
+        (paths.CACHE / clip.name).mkdir(parents=True, exist_ok=True)
+        for m in todo:
+            cp = cache_paths(clip.name, m, res)
+            Image.fromarray(frames[anchor_frame(m)]).save(cp["png"])
+            big = _load_anchor(cp["png"]) if res == "2048" else np.ascontiguousarray(frames[anchor_frame(m)])
+            with torch.inference_mode():
+                lat = encode_video_condition(vae, _prepare_pixels(torch.from_numpy(big).unsqueeze(0)).to(dev, torch.float16))[0].cpu().float()
+            atomic_save(cp["latent"], {"latents_ref_000_image": lat})
+            caption = _canonical_caption(caption_path(clip, m).read_text(), f"{clip.name}/static_c{m:02d}")
+            hidden, tags = static_text_states(proc, model, big, cp["png"], clip / "conditions_coarse", anchor_frame(m), caption)
+            atomic_save(cp["text"], {"hidden": hidden.to("cpu", torch.bfloat16), "tags": tags.to("cpu", torch.int64)},
+                        {"system_prompt": "w0", "duv": "start frame repeated", "start_chunk": str(m)})
+            gc.collect(); torch.cuda.empty_cache()
+        done += len(todo)
+        print("static_text", clip.name, len(todo), "start frames", flush=True)
+    return done
+
+
+def static_text_states(proc, model, picture1: np.ndarray, picture1_png: Path, cond: Path, frame: int, caption: str):
+    """Qwen states [L,5120] and token tags of a static caption: Qwen sees Picture 1 ([H,W,3] uint8, saved at
+    picture1_png) and, as the Ref2VA video reference, the proxy of `frame` repeated -- nothing of later
+    frames -- with the first-window system prompt."""
+    from musubi_tuner.minimax_h3.media import H3AudioSource, H3Record, H3Reference
+    from musubi_tuner.minimax_h3.text_encoder import H3TextVisual, build_presentation, encode_h3_presentation
+    from cwm_h3_inference.duv import load_duv_frame
+    anchor = torch.from_numpy(np.ascontiguousarray(picture1)).unsqueeze(0)
+    f0 = load_duv_frame(cond, frame, for_qwen=True)
+    duv = torch.from_numpy(np.ascontiguousarray(np.stack([f0] * len(QWEN_OFFSETS))))
+    video_key = cond / f"{frame:06d}.semantic_id.png"
+    placeholder = cond / ".h3-presentation-placeholder"
+    record = H3Record(video_path=placeholder, caption=caption, target_audio=H3AudioSource(path=placeholder, embedded=False),
+                      references=(H3Reference(type="image", path=picture1_png),
+                                  H3Reference(type="video", path=video_key, duration_seconds=WINDOW_FRAMES / FPS)),
+                      jsonl_line=0)
+    plain = build_presentation(record, "ref2va", {picture1_png: H3TextVisual(anchor),
+                                                  video_key: H3TextVisual(duv, tuple(o / FPS for o in QWEN_OFFSETS))})
+    with torch.inference_mode():
+        return encode_h3_presentation(proc, model, build_ref2va_presentation(proc, plain, 0, caption))
+
+
+def _job_running(job: str) -> bool:
+    import subprocess
+    return bool(subprocess.run(["squeue", "-h", "-j", job], capture_output=True, text=True).stdout.strip())
+
+
+def text_done(clip: Path) -> bool:
+    return all((paths.CACHE / clip.name / f"w{w:02d}.text.safetensors").exists() for w in range(windows_of(clip)))
+
+
+def select_clips(shard: bool, reverse: bool, only_incomplete: bool):
+    clips = [d for d in sorted(paths.DATASET.glob("clip_*")) if d.is_dir() and qc_pass(d) and has_prompts(d)]
+    if only_incomplete:
+        clips = [c for c in clips if not text_done(c)]
+    if reverse:
+        clips = clips[::-1]
+    if shard:
+        i, n = int(os.environ["SLURM_PROCID"]), int(os.environ["SLURM_NTASKS"])
+        clips = clips[i::n]
+    return clips
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--stage", choices=["visual", "text", "both", "stream", "static_text"], required=True)
+    ap.add_argument("clips", nargs="*", type=Path)
+    ap.add_argument("--shard", action="store_true", help="take this SLURM task's share of the dataset")
+    ap.add_argument("--reverse", action="store_true")
+    ap.add_argument("--only-incomplete", action="store_true", help="skip clips whose text stage is done")
+    ap.add_argument("--follow-job", default="", help="static_text: keep encoding new captions while this SLURM job runs")
+    ap.add_argument("--picture1", choices=["native", "2048"], default="native", help="static_text: Picture 1 resolution")
+    a = ap.parse_args()
+    clips = [c.resolve() for c in a.clips] or select_clips(a.shard, a.reverse, a.only_incomplete)
+    print(f"{len(clips)} clips", flush=True)
+    if a.stage in ("visual", "both"):
+        stage_visual(clips)
+    if a.stage in ("text", "both"):
+        stage_text(clips)
+    if a.stage == "stream":
+        stage_stream(clips)
+    if a.stage == "static_text":
+        stage_static_text(clips, a.follow_job, a.picture1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rtcwm/train/train.py b/rtcwm/train/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..99d7e0557644becb84bd5a72a09bcdc8a900d1b9
--- /dev/null
+++ b/rtcwm/train/train.py
@@ -0,0 +1,156 @@
+"""Causal (streaming) LoRA training of the CWM, initialised from the released CWM LoRA.
+
+    torchrun --nnodes N --nproc-per-node 8 ... -m rtcwm.train.train --out DIR [--steps 150 | --epochs E]
+
+* diffusion forcing on 6-chunk spans cut from the clips' latent streams, each with a drawn start frame
+  (rtcwm.train.dataset, rtcwm.model.objectives)
+* the frozen H3 blocks are FSDP-sharded over the 8 GPUs of each node; the LoRA weights are replicated
+  and their gradients all-reduced asynchronously as each is produced in backward
+* one window per GPU per step; steps = epochs * samples / world unless --steps is given
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import time
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from .. import paths
+from ..model import loading
+from ..model.layout import SPAN_CHUNKS
+from ..model.objectives import LOSSES, OBJECTIVES
+from .dataset import StartFrames, list_samples, load_window, silence
+
+
+def parse():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--objective", choices=OBJECTIVES, default="df")
+    ap.add_argument("--init-lora", type=Path, default=paths.CWM_LORA)
+    ap.add_argument("--steps", type=int, default=0, help="0 = epochs * samples / world")
+    ap.add_argument("--epochs", type=float, default=2.0)
+    ap.add_argument("--lr", type=float, default=2e-5)
+    ap.add_argument("--lr-min", type=float, default=1e-6)
+    ap.add_argument("--warmup", type=int, default=10)
+    ap.add_argument("--save-every", type=int, default=10)
+    ap.add_argument("--val-clips", type=int, default=12)
+    ap.add_argument("--max-samples", type=int, default=0, help="smoke tests: cap the training set")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--picture1", choices=["native", "2048"], default="native", help="Picture 1 resolution")
+    ap.add_argument("--attn", choices=["fa3", "flex"], default="fa3", help="the same mask: FlashAttention-3 prefix calls or flex")
+    ap.add_argument("--level-shift", type=float, default=3.0, help="DF: flow shift of the per-chunk noise levels")
+    ap.add_argument("--act-offload", choices=["on", "off"], default="off", help="checkpointed block inputs on the CPU")
+    return ap.parse_args()
+
+
+def node_group(rank: int, world: int):
+    """One FSDP shard group per node (NVLink); every rank must create every group."""
+    lw = int(os.environ.get("LOCAL_WORLD_SIZE", torch.cuda.device_count()))
+    groups = [dist.new_group(list(range(n * lw, (n + 1) * lw))) for n in range(world // lw)]
+    return groups[rank // lw]
+
+
+def resume_point(a):
+    """(steps done, LoRA to start from): the latest checkpoint in --out (a requeued or resubmitted run --
+    lowprio jobs are preempted and requeued -- continues from it), else --init-lora."""
+    done = sorted(int(p.stem.split("step")[1]) for p in a.out.glob("lora_step*.safetensors"))
+    return (done[-1], a.out / f"lora_step{done[-1]:05d}.safetensors") if done else (0, a.init_lora)
+
+
+def cosine_with_warmup(opt, a, steps):
+    return torch.optim.lr_scheduler.LambdaLR(opt, lambda i: min(1.0, (i + 1) / a.warmup) * (
+        a.lr_min / a.lr + (1 - a.lr_min / a.lr) * 0.5 * (1 + math.cos(math.pi * min(i, steps) / steps))))
+
+
+def main():
+    a = parse()
+    dist.init_process_group("nccl")
+    rank, world = dist.get_rank(), dist.get_world_size()
+    local = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(local)
+    dev = torch.device("cuda", local)
+    log = (lambda *m: print(f"[{time.strftime('%H:%M:%S')}]", *m, flush=True)) if rank == 0 else (lambda *m: None)
+    group = node_group(rank, world)
+
+    train, val, dropped = list_samples(a.val_clips, a.seed, a.picture1)
+    if a.max_samples:
+        train = train[:a.max_samples]
+    steps = a.steps or int(math.ceil(a.epochs * len(train) / world))
+    log(f"objective={a.objective} span_chunks={SPAN_CHUNKS} level_shift={a.level_shift} attn={a.attn} act_offload={a.act_offload} world={world} train={len(train)} val={len(val)} dropped_flagged={dropped} steps={steps}")
+    if rank == 0:
+        a.out.mkdir(parents=True, exist_ok=True)
+        (a.out / "val_samples.json").write_text(json.dumps(val))
+        (a.out / "args.json").write_text(json.dumps({k: str(v) for k, v in vars(a).items()}, indent=1))
+
+    tr = loading.load_base()
+    start, init = resume_point(a)
+    if start:
+        log(f"resuming after step {start}: {init.name} (optimizer moments start afresh)")
+    net = loading.attach_trainable_lora(tr, init)
+    loading.shard_blocks(tr, group, dev)
+    net.to(dev, torch.float32)
+    tr.train()
+    params = list(net.parameters())
+    handles = []
+
+    def reduce_hook(p):
+        handles.append(dist.all_reduce(p.grad, op=dist.ReduceOp.AVG, async_op=True))
+    for p in params:
+        p.requires_grad_(True)
+        p.register_post_accumulate_grad_hook(reduce_hook)
+    log(f"LoRA trainable={sum(p.numel() for p in params) / 1e6:.1f}M mem_after_shard={torch.cuda.memory_allocated() / 2**30:.1f}G")
+    opt = torch.optim.AdamW(params, lr=a.lr, betas=(0.9, 0.999), weight_decay=0.01)
+    sched = cosine_with_warmup(opt, a, steps)
+    for _ in range(start):                       # the schedule where the resumed run left it
+        sched.step()
+    loss_fn, audio = LOSSES[a.objective], silence()
+    offload = a.act_offload == "on"
+    gen = torch.Generator().manual_seed(a.seed * 1000 + rank)
+    starts = StartFrames(train, a.seed * 1000 + rank, a.picture1)
+
+    order, ptr, mask_cache, t0 = [], 0, {}, time.time()
+    for it in range(steps):
+        if ptr >= len(order):
+            order = train[:]
+            random.Random(a.seed + it).shuffle(order)
+            order, ptr = order[rank::world], 0
+        clip, j = order[ptr]
+        ptr += 1
+        j, m = starts.draw(clip, j)
+        if it < start:                           # resumed: replay the sample order up to where it stopped
+            continue
+        loss = loss_fn(tr, load_window(clip, j, audio, start=m, res=a.picture1), dev, gen, mask_cache, offload=offload,
+                       attn=a.attn, level_shift=a.level_shift)
+        loss.backward()
+        for h in handles:
+            h.wait()
+        handles.clear()
+        gn = torch.nn.utils.clip_grad_norm_(params, 1.0)
+        opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
+        stat = loss.detach().reshape(1)
+        dist.all_reduce(stat, op=dist.ReduceOp.AVG)
+        if rank == 0:
+            el = time.time() - t0
+            log(f"step {it + 1}/{steps} loss_v={stat[0]:.4f} gnorm={float(gn):.3f} lr={sched.get_last_lr()[0]:.2e} "
+                f"mem={torch.cuda.max_memory_allocated() / 2**30:.0f}G {el / (it + 1 - start):.1f}s/step "
+                f"eta={(steps - it - 1) * el / (it + 1 - start) / 3600:.2f}h")
+            with open(a.out / "train_log.jsonl", "a") as fh:
+                fh.write(json.dumps(dict(step=it + 1, loss_v=float(stat[0]), gnorm=float(gn), lr=sched.get_last_lr()[0],
+                                         sec=el)) + "\n")
+            if (it + 1) % a.save_every == 0 or it + 1 == steps:
+                p = a.out / f"lora_step{it + 1:05d}.safetensors"
+                net.save_weights(str(p), torch.bfloat16, {"base": "cwm", "step": str(it + 1), "objective": a.objective,
+                                                    "span_chunks": str(SPAN_CHUNKS), "level_shift": str(a.level_shift)})
+                log("saved", p)
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()