prep stage

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)

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()

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"

cwm_h3_inference/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

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

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())

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

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"

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])

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

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

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)

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.

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.

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

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())

musubi_tuner/minimax_h3/__init__.py

@@ -0,0 +1 @@
+"""MiniMax-H3 model, cache, training, and sampling support."""

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,
+    )

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

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

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,
+    )

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),
+    )

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,
+    )

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()}"

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,
+    )

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)

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

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)

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()

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]

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

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

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()

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

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

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

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

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

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"])
+
+"""

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

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

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"

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."""

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")

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()

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)

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