Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/man_laughing_audio.wav filter=lfs diff=lfs merge=lfs -text
+examples/man_laughing_video.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/man_smiling_studio_audio.wav filter=lfs diff=lfs merge=lfs -text
+examples/man_smiling_studio_video.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/woman_laughing_studio_audio.wav filter=lfs diff=lfs merge=lfs -text
+examples/woman_laughing_studio_video.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,61 @@
 ---
-title: Viddub
-emoji: 😻
-colorFrom: yellow
-colorTo: green
+title: VidDub
+emoji: 🎬
+colorFrom: green
+colorTo: red
 sdk: gradio
 sdk_version: 6.31.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+python_version: "3.10"
+short_description: Two-step video dubbing — re-sync lips to new speech
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# VidDub — Replacement-Aware Distillation for Two-Step Video Dubbing
+
+**VidDub** ([paper](https://huggingface.co/papers/2611.03404)) is a two-step
+diffusion video-dubbing editor: given a source talking-head video and new
+target speech, it regenerates the speaker's lower face so the lips match the
+new audio — everything else is kept from the source. It is trained on a
+Wan2.1-VACE-1.3B backbone with a Whisper-small audio adapter, distilled with
+DMD to just **2 denoising steps** (flow-euler, shift 5, audio CFG 1).
+
+- 💻 Code: [LAW1223/VidDub](https://github.com/LAW1223/VidDub)
+- 🤗 Weights: [AIPeanutman/VidDub](https://huggingface.co/AIPeanutman/VidDub)
+  (base: [Wan-AI/Wan2.1-VACE-1.3B](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B),
+  audio encoder: [openai/whisper-small](https://huggingface.co/openai/whisper-small))
+
+## Usage
+
+1. Upload a source video with a single, mostly front-facing face
+   (≥ 81 frames, i.e. ~3.4 s; longer videos are processed in 81-frame windows
+   and truncated at the frame cap).
+2. Upload the target speech audio (any common format).
+3. Press **Dub video**. The output keeps the source resolution and frame rate
+   and carries the new speech.
+
+The two-step DMD checkpoint is fixed to 2 steps / audio CFG 1 — the sampler
+has no user-tunable steps. Advanced options: noise seed, the paper's
+lower-face color harmonization, mask feather, and the frame cap.
+
+## Implementation notes
+
+The app replicates the reference implementation (`scripts/infer.py`) 1:1:
+MediaPipe FaceMesh landmarks → 512² face crop + lower-face mask → VACE latents
+(x0 / inactive / ref) + empty-prompt T5 context + Whisper features →
+two-step bidirectional rollout → Lab color harmonization → feathered paste-back
+into the original frames → AAC muxing at 48 kHz stereo.
+
+One deviation: dubbing never uses a text prompt, so the T5 context of the
+empty string is precomputed once (18 KB tensor in `assets/empty_context.pt`)
+instead of downloading the 11 GB umt5-xxl encoder at startup.
+
+## Licenses / credits
+
+- Example videos: CC0, from
+  [linoyts/repo-to-space-example-videos](https://huggingface.co/datasets/linoyts/repo-to-space-example-videos).
+- Example speech: CC BY 4.0, from
+  [LibriSpeech](https://huggingface.co/datasets/openslr/librispeech_asr)
+  (Vassil Panayotov / ILSVRC & ISLR, 2015).
+- VidDub model weights are released under the repository's own
+  source-available license; Wan2.1 base weights under Apache-2.0.

app.py

@@ -0,0 +1,337 @@
+"""VidDub: two-step replacement-aware video dubbing.
+
+Upload a source video (a talking head) and target speech audio; VidDub
+regenerates the lower face so the lips match the new audio in 2 denoising
+steps, pasting the result back into the original video.
+
+Reference implementation: https://github.com/LAW1223/VidDub (scripts/infer.py).
+Model: AIPeanutman/VidDub on top of Wan2.1-VACE-1.3B + whisper-small.
+"""
+import json
+import math
+import os
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import spaces  # MUST be imported before torch / anything CUDA-touching
+
+# The Space repo bundles the VidDub source (viddub/ + data/) at /home/user/app.
+APP_ROOT = Path(__file__).resolve().parent
+for p in (str(APP_ROOT),):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import decord
+import gradio as gr
+import librosa
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from safetensors.torch import load_file
+
+from viddub.causal.model import CausalConfig, CausalWanDub
+from viddub.causal.runtime import autocast
+from viddub.causal.sampling import rollout, schedule
+from viddub.foundation.pipelines.wan_video import ModelConfig, WanVideoPipeline
+from viddub.full_context import runtime as full_context_runtime
+from viddub.full_context.audio import DubAudioAdapter, WhisperFeatures
+from viddub.full_context.inference import dub_video
+from viddub.full_context.runtime import FPS, SR, ZeroLatents, build_inputs, write_mp4
+
+from data.s02_make_clips import NUM_FRAMES
+
+DEVICE = "cuda"
+VIDDUB_REPO = "AIPeanutman/VidDub"
+BASE_REPO = "Wan-AI/Wan2.1-VACE-1.3B"
+WHISPER_REPO = "openai/whisper-small"
+
+MAX_FRAMES = 242  # 3 windows x 81 frames (stride 64) + tail; ~10s @ 24fps
+
+
+# ---------------------------------------------------------------- weight prep
+def prepare_weights():
+    """Download every checkpoint the runtime needs into one local root."""
+    root = Path(os.environ.get("VIDDUB_WEIGHT_ROOT", APP_ROOT / "weights"))
+    root.mkdir(parents=True, exist_ok=True)
+    vace_dir = root / "Wan2.1-VACE-1.3B"
+    need = {
+        "diffusion": vace_dir / "diffusion_pytorch_model.safetensors",
+        "vae": vace_dir / "Wan2.1_VAE.pth",
+    }
+    if not all(p.is_file() for p in need.values()):
+        snapshot_download(
+            BASE_REPO, repo_type="model",
+            allow_patterns=["diffusion_pytorch_model.safetensors", "Wan2.1_VAE.pth"],
+            local_dir=str(vace_dir))
+    whisper_dir = root / "whisper-small"
+    if not (whisper_dir / "model.safetensors").is_file():
+        snapshot_download(
+            WHISPER_REPO, repo_type="model",
+            allow_patterns=["*.json", "*.txt", "model.safetensors", "vocab.json",
+                            "merges.txt", "preprocessor_config.json",
+                            "tokenizer_config.json", "normalizer.json",
+                            "special_tokens_map.json", "added_tokens.json"],
+            local_dir=str(whisper_dir))
+    release_dir = root / "VidDub"
+    if not (release_dir / "config.json").is_file():
+        snapshot_download(
+            VIDDUB_REPO, repo_type="model",
+            allow_patterns=["*.json", "dit/*", "vace/*", "audio_adapter/*"],
+            local_dir=str(release_dir))
+    full_context_runtime.WEIGHTS = root
+    full_context_runtime.VACE_DIR = vace_dir
+    return release_dir
+
+
+def load_release(model_dir: Path):
+    """Load the VidDub release exactly as scripts/infer.py does."""
+    config = json.loads((model_dir / "config.json").read_text())
+    if config.get("model_type") != "viddub":
+        raise ValueError("The model directory is not a VidDub release")
+    if float(config["inference_audio_cfg"]) != 1.0:
+        raise ValueError("Published VidDub checkpoints require inference audio CFG 1")
+    if config["stage"] != "bidirectional_sft" and (
+        config["stage"] != "dmd"
+        or config.get("paper_selected") is not True
+        or config.get("causal") is not False
+        or config.get("source_run") != "FULL2_DMD_R4"):
+        raise ValueError("The released two-step path requires FULL2_DMD_R4")
+
+    # Same T5-skip path as viddub/causal/runtime.py: with a precomputed
+    # empty-prompt context, the 11 GB T5 encoder is never loaded.
+    empty_context = torch.load(os.environ["VIDDUB_EMPTY_CONTEXT"],
+                               map_location="cpu", weights_only=True)
+    if not (torch.is_tensor(empty_context) and empty_context.ndim == 2
+            and empty_context.shape[1] == 4096 and 1 <= empty_context.shape[0] <= 512
+            and torch.isfinite(empty_context).all()):
+        raise ValueError("assets/empty_context.pt must be a finite [L,4096] tensor")
+    pipe = WanVideoPipeline.from_pretrained(
+        torch_dtype=torch.bfloat16, device=DEVICE,
+        model_configs=[ModelConfig(path=str(full_context_runtime.VACE_DIR
+                                            / "diffusion_pytorch_model.safetensors")),
+                       ModelConfig(path=str(full_context_runtime.VACE_DIR
+                                            / "Wan2.1_VAE.pth"))],
+        tokenizer_config=None)
+    for m in (pipe.dit, pipe.vace, pipe.vae):
+        m.requires_grad_(False).eval()
+    pipe._wan_dub_empty_context = empty_context.to(pipe.device, pipe.torch_dtype)
+    whisper = WhisperFeatures(str(full_context_runtime.WEIGHTS
+                                  / "whisper-small")).to(DEVICE, torch.bfloat16)
+    paths = {name: model_dir / rel for name, rel in config["components"].items()}
+    missing = [str(p) for p in paths.values() if not p.is_file()]
+    if missing:
+        raise FileNotFoundError(f"Missing model files: {missing}")
+    # The VidDub release ships bf16 DiT/VACE and fp32 audio adapter; the
+    # reference loader loads them straight onto the device with strict=True.
+    pipe.dit.load_state_dict(load_file(str(paths["dit"])), strict=True)
+    pipe.vace.load_state_dict(load_file(str(paths["vace"])), strict=True)
+    adapter = DubAudioAdapter().to(DEVICE)
+    adapter.load_state_dict(load_file(str(paths["audio_adapter"])), strict=True)
+    adapter.eval()
+    return config, pipe, whisper, adapter
+
+
+def two_step_sampler(model, config):
+    """Two-step flow_euler sampler with a full-context (bidirectional) rollout."""
+    @torch.no_grad()
+    def sample(pipe, whisper, adapter, zero, clip, audio, steps=None,
+               audio_cfg=1.0, seed=0, fps=FPS):
+        trained_steps = int(config["student_steps"])
+        num_frames = len(clip["frames"])
+        video = full_context_runtime.to_video_tensor(clip["frames"], pipe)
+        keep = torch.from_numpy(~clip["mask"]).to(video)[None, None, None]
+        latents = {
+            "x0": pipe.vae.encode(video, device=pipe.device)[0],
+            "inactive": pipe.vae.encode(video * keep, device=pipe.device)[0],
+            "ref": pipe.vae.encode(
+                full_context_runtime.to_video_tensor(clip["ref"][None], pipe),
+                device=pipe.device)[0],
+        }
+        x0, vace_context, context = build_inputs(
+            latents, clip["mask"], num_frames, zero,
+            drop_ref=int(config["num_reference_latents"]) == 0)
+        whisper_hs = whisper([audio], pipe.device, torch.bfloat16)
+        model.config = CausalConfig(
+            int(config["latent_chunk_size"]),
+            int(config["num_reference_latents"]), float(fps))
+        conditions = {
+            "context": context.to(torch.bfloat16),
+            "vace_context": vace_context.to(torch.bfloat16),
+            "whisper_hs": whisper_hs.to(torch.bfloat16),
+            "num_frames": num_frames,
+        }
+        noise = torch.randn(
+            x0.shape, generator=torch.Generator(device=x0.device).manual_seed(seed),
+            device=x0.device, dtype=torch.bfloat16)
+        with autocast(x0.device, "bf16"):
+            output = rollout(
+                model, x0.to(torch.bfloat16), conditions,
+                schedule(trained_steps, float(config["shift"]), x0.device),
+                noise=noise, use_cache=bool(config.get("causal", False)),
+                audio_cfg=1.0)
+        video = pipe.vae.decode(
+            output[:, :, int(config["num_reference_latents"]):].to(pipe.torch_dtype),
+            device=pipe.device)
+        return ((video[0].permute(1, 2, 3, 0).float().clamp(-1, 1) + 1)
+                * 127.5).round().byte().cpu().numpy()
+
+    return sample
+
+
+# ------------------------------------------------------------ module-scope load
+# Dubbing is audio-driven: the only text condition is the empty prompt. The
+# bundled assets/empty_context.pt holds the precomputed T5 embedding of "",
+# so the 11 GB umt5-xxl encoder is never downloaded or loaded (this is the
+# repo's own VIDDUB_SKIP_TEXT_ENCODER path, see viddub/causal/runtime.py).
+os.environ["VIDDUB_SKIP_TEXT_ENCODER"] = "1"
+os.environ["VIDDUB_EMPTY_CONTEXT"] = str(APP_ROOT / "assets" / "empty_context.pt")
+
+print("Downloading weights ...", flush=True)
+RELEASE_DIR = prepare_weights()
+print("Loading VidDub runtime ...", flush=True)
+CONFIG, PIPE, WHISPER, ADAPTER = load_release(RELEASE_DIR)
+ZERO = ZeroLatents(PIPE)
+INFERENCE_MODEL = CausalWanDub(
+    PIPE.dit, PIPE.vace, ADAPTER,
+    CausalConfig(int(CONFIG["latent_chunk_size"]),
+                 int(CONFIG["num_reference_latents"]), float(CONFIG["fps"])),
+    causal=False).eval().requires_grad_(False)
+SAMPLE_FN = two_step_sampler(INFERENCE_MODEL, CONFIG)
+RUNTIME = {"models": (PIPE, WHISPER, ADAPTER, ZERO), "config": CONFIG,
+           "steps": int(CONFIG["student_steps"]), "sample_fn": SAMPLE_FN}
+print("VidDub runtime ready.", flush=True)
+
+
+# ----------------------------------------------------------------- inference
+def _probe_duration(video_path: str) -> float:
+    """Estimate the GPU duration from the input video's frame count."""
+    try:
+        n = len(decord.VideoReader(video_path, num_threads=1))
+    except Exception:
+        n = NUM_FRAMES
+    n = max(n, NUM_FRAMES)
+    n_windows = max(1, math.ceil((n - NUM_FRAMES) / 64) + 1)
+    # 81-frame clip: ~11 s on the RTX PRO 6000; per extra window ~10 s more,
+    # plus VAE encode/decode and MediaPipe preprocessing. Measured, +15% margin.
+    return min(420, 45 + 11 * n_windows)
+
+
+@spaces.GPU(duration=_probe_duration)
+def dub(video: str, audio: str, seed: int = 0, harmonize: bool = True,
+        feather: int = 15, max_frames: int = MAX_FRAMES,
+        progress=gr.Progress(track_tqdm=False)):
+    """Re-dub a talking-head video with new speech audio.
+
+    Args:
+        video: source video with a mostly front-facing face (mp4/webm).
+        audio: target speech (any format; resampled to 16 kHz mono).
+        seed: noise seed for the two-step sampler.
+        harmonize: apply the paper's lower-face color harmonization.
+        feather: mask feather radius in crop pixels.
+        max_frames: cap on processed frames (longer videos are truncated).
+
+    Returns:
+        The dubbed video with the original resolution/fps and the new speech.
+    """
+    if not video or not audio:
+        raise gr.Error("Please provide both a source video and target audio.")
+    t0 = time.perf_counter()
+    video = str(video)
+    audio = str(audio)
+    reader = decord.VideoReader(video, num_threads=2)
+    fps = float(reader.get_avg_fps())
+    frames = reader[:].asnumpy()
+    if len(frames) > max_frames:
+        frames = frames[:max_frames]
+    if len(frames) < NUM_FRAMES:
+        raise gr.Error(
+            f"Source video needs at least {NUM_FRAMES} frames "
+            f"({NUM_FRAMES / FPS:.1f} s); got {len(frames)}.")
+    waveform = librosa.load(audio, sr=SR, mono=True)[0]
+    audio_samples = round(len(frames) / fps * SR)
+    waveform = np.pad(
+        waveform[:audio_samples], (0, max(0, audio_samples - len(waveform))))
+    progress((0.1,), desc="Dubbing ...")
+    generated = dub_video(
+        RUNTIME["models"], frames, waveform,
+        steps=RUNTIME["steps"], audio_cfg=1.0, feather=feather, seed=seed,
+        fps=fps, overlap=17, mode="auto", harmonize=harmonize,
+        sample_fn=RUNTIME["sample_fn"],
+        log=lambda m: print(m, flush=True))
+    out_path = Path(tempfile.mkdtemp()) / "dubbed.mp4"
+    write_mp4(out_path, generated, waveform, fps=fps)
+    print(f"dub finished in {time.perf_counter() - t0:.1f}s", flush=True)
+    return str(out_path)
+
+
+# ----------------------------------------------------------------------- UI
+EXAMPLES = [
+    ["examples/man_laughing_video.mp4", "examples/man_laughing_audio.wav", 0],
+    ["examples/woman_laughing_studio_video.mp4", "examples/woman_laughing_studio_audio.wav", 0],
+    ["examples/man_smiling_studio_video.mp4", "examples/man_smiling_studio_audio.wav", 0],
+]
+
+with gr.Blocks(theme=gr.themes.Citrus(), title="VidDub") as demo:
+    gr.Markdown(
+        """
+        # 🎬 VidDub — Replacement-Aware Distillation for Two-Step Video Dubbing
+
+        **VidDub** takes a source talking-head video and a new speech clip and
+        re-renders the speaker's lower face so the lips match the new audio —
+        in **2 denoising steps** on a Wan2.1-VACE-1.3B backbone.
+
+        📄 [Paper](https://huggingface.co/papers/2611.03404) ·
+        💻 [GitHub](https://github.com/LAW1223/VidDub) ·
+        🤗 [AIPeanutman/VidDub](https://huggingface.co/AIPeanutman/VidDub)
+
+        Upload a short clip with a single, mostly front-facing face and the
+        audio you want the face to speak. Everything except the mouth region is
+        kept from the source video.
+        """)
+    with gr.Row():
+        with gr.Column():
+            video_in = gr.Video(label="Source video (talking head)", sources=["upload"])
+            audio_in = gr.Audio(label="Target speech audio", sources=["upload"], type="filepath")
+            run_btn = gr.Button("Dub video", variant="primary")
+        with gr.Column():
+            video_out = gr.Video(label="Dubbed video", autoplay=True)
+    with gr.Accordion("Advanced options", open=False):
+        seed = gr.Number(label="Seed", value=0, precision=0)
+        harmonize = gr.Checkbox(label="Color harmonization (paper default)", value=True)
+        feather = gr.Slider(1, 40, value=15, step=1, label="Mask feather (px)")
+        max_frames = gr.Slider(NUM_FRAMES, MAX_FRAMES, value=MAX_FRAMES, step=1,
+                               label="Max frames (longer videos are truncated)")
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[video_in, audio_in, seed],
+        fn=dub,
+        outputs=[video_out],
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+    run_btn.click(
+        fn=dub,
+        inputs=[video_in, audio_in, seed, harmonize, feather, max_frames],
+        outputs=[video_out],
+        api_name="dub",
+    )
+    gr.Markdown(
+        """
+        ---
+        **Notes.** The two-step DMD checkpoint is fixed: 2 denoising steps,
+        audio CFG 1, 81-frame windows with 17-frame overlap. Videos need a
+        detectable, mostly front-facing face and at least 81 frames
+        (~3.4 s). VidDub-Bench is CC BY-NC; the bundled example videos come
+        from the CC0 pool at
+        [linoyts/repo-to-space-example-videos](https://huggingface.co/datasets/linoyts/repo-to-space-example-videos)
+        and the speech from
+        [LibriSpeech](https://huggingface.co/datasets/openslr/librispeech_asr)
+        (CC BY 4.0).
+        """)
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

assets/empty_context.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad9140ae8591e91b64c3d5dbeb183c6701f19d419c8b5f43fe364146042e88ae
+size 18003

data/__init__.py

@@ -0,0 +1 @@
+"""VidDub data preparation and preprocessing tools."""

data/common.py

@@ -0,0 +1,38 @@
+"""Paths and JSONL helpers shared by the public data pipeline."""
+import json
+import os
+from pathlib import Path
+
+PROJECT = Path(__file__).resolve().parents[1]
+DATA = Path(os.environ.get("VIDDUB_DATA_ROOT", PROJECT / "data"))
+MANIFEST = DATA / "manifest"
+WEIGHTS = Path(os.environ.get("VIDDUB_WEIGHT_ROOT", PROJECT / "weights"))
+
+
+def read_jsonl(path):
+    rows = []
+    with Path(path).open() as handle:
+        for line in handle:
+            if line.strip():
+                rows.append(json.loads(line))
+    return rows
+
+
+def write_jsonl(path, rows):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    with tmp.open("w") as handle:
+        for row in rows:
+            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+    tmp.replace(path)
+
+
+def resolve_video_path(row):
+    """Resolve artifact-relative media paths against VIDDUB_DATA_ROOT."""
+    path = Path(row["video_path"]).expanduser()
+    return path if path.is_absolute() else DATA / path
+
+
+def shard(rows, shard_id, num_shards):
+    return rows[shard_id::num_shards]

data/s01_analyze.py

@@ -0,0 +1,182 @@
+"""Per video: face landmarks (mediapipe), 16k mono wav, SyncNet offset/confidence.
+
+Outputs:
+  data/face/<sub>/<id>.npz            lm [T,468,2] float16 (nan = no face), nfaces [T]
+  data/audio/<sub>/<id>.wav           16kHz mono
+  data/manifest/s01/shard_XXXX.jsonl  per-video stats (append; resumable)
+"""
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import cv2
+import decord
+import mediapipe as mp
+import numpy as np
+import python_speech_features
+import torch
+from scipy import signal
+from scipy.io import wavfile
+
+from .common import DATA, MANIFEST, PROJECT, WEIGHTS, read_jsonl, resolve_video_path, shard
+
+LATENTSYNC_ROOT = Path(os.environ.get("VIDDUB_LATENTSYNC_ROOT", WEIGHTS / "LatentSync"))
+
+SYNCNET_PATH = Path(os.environ.get(
+    "VIDDUB_SYNCNET_PATH", WEIGHTS / "LatentSync-1.5/auxiliary/syncnet_v2.model"))
+LM_NOSE, LM_CHEEK_R, LM_CHEEK_L, LM_LIP_UP, LM_LIP_LOW, LM_FOREHEAD, LM_CHIN = 1, 234, 454, 13, 14, 10, 152
+
+
+def sub_dir(vid):
+    return vid[-6:-3]
+
+
+def load_syncnet(device):
+    syncnet_root = LATENTSYNC_ROOT / "eval/syncnet"
+    if str(syncnet_root) not in sys.path:
+        sys.path.insert(0, str(syncnet_root))
+    from syncnet import S
+    net = S(num_layers_in_fc_layers=1024).to(device).eval()
+    net.load_state_dict(torch.load(SYNCNET_PATH, map_location="cpu", weights_only=True))
+    return net
+
+
+def detect_landmarks(
+    frames,
+    min_detection_confidence=0.5,
+    min_tracking_confidence=0.5,
+):
+    T, H, W = frames.shape[:3]
+    lm = np.full((T, 468, 2), np.nan, dtype=np.float32)
+    nfaces = np.zeros(T, dtype=np.int8)
+    with mp.solutions.face_mesh.FaceMesh(
+        static_image_mode=False,
+        max_num_faces=2,
+        min_detection_confidence=min_detection_confidence,
+        min_tracking_confidence=min_tracking_confidence,
+    ) as face_mesh:
+        for t in range(T):
+            res = face_mesh.process(frames[t])
+            if not res.multi_face_landmarks:
+                continue
+            nfaces[t] = len(res.multi_face_landmarks)
+            faces = [np.array([(p.x * W, p.y * H) for p in f.landmark[:468]]) for f in res.multi_face_landmarks]
+            lm[t] = max(faces, key=lambda a: np.ptp(a[:, 1]))  # largest face
+    return lm, nfaces
+
+
+def syncnet_crops(frames, lm, crop_scale=0.4):
+    """Replicates syncnet_python crop (S3FD box -> landmark bbox), median-filtered track."""
+    valid = ~np.isnan(lm[:, 0, 0])
+    idx = np.arange(len(lm))
+    xs = np.interp(idx, idx[valid], (np.nanmin(lm[valid, :, 0], 1) + np.nanmax(lm[valid, :, 0], 1)) / 2)
+    ys = np.interp(idx, idx[valid], (np.nanmin(lm[valid, :, 1], 1) + np.nanmax(lm[valid, :, 1], 1)) / 2)
+    ss = np.interp(idx, idx[valid], np.maximum(np.ptp(lm[valid, :, 0], 1), np.ptp(lm[valid, :, 1], 1)) / 2)
+    xs, ys, ss = (signal.medfilt(v, kernel_size=13) for v in (xs, ys, ss))
+    crops = []
+    for t in range(len(frames)):
+        bs = ss[t]
+        bsi = int(bs * (1 + 2 * crop_scale))
+        img = np.pad(frames[t], ((bsi, bsi), (bsi, bsi), (0, 0)), "constant", constant_values=110)
+        my, mx = ys[t] + bsi, xs[t] + bsi
+        face = img[int(my - bs):int(my + bs * (1 + 2 * crop_scale)), int(mx - bs * (1 + crop_scale)):int(mx + bs * (1 + crop_scale))]
+        crops.append(cv2.resize(face, (224, 224)))
+    return np.stack(crops)
+
+
+@torch.no_grad()
+def syncnet_score(net, crops_rgb_25fps, audio_int16, device, vshift=15, batch=64):
+    im = crops_rgb_25fps[..., ::-1].astype(np.float32)  # syncnet_python feeds cv2 BGR frames
+    im = torch.from_numpy(np.ascontiguousarray(im)).permute(3, 0, 1, 2)[None]  # 1,3,T,224,224
+    mfcc = np.stack(python_speech_features.mfcc(audio_int16, 16000), 1)  # 13, T_a
+    cc = torch.from_numpy(mfcc.astype(np.float32))[None, None]
+    n = min(im.shape[2], len(audio_int16) // 640) - 5
+    im_feat, cc_feat = [], []
+    for i in range(0, n, batch):
+        r = range(i, min(n, i + batch))
+        im_feat.append(net.forward_lip(torch.cat([im[:, :, v:v + 5] for v in r]).to(device)).cpu())
+        cc_feat.append(net.forward_aud(torch.cat([cc[:, :, :, v * 4:v * 4 + 20] for v in r]).to(device)).cpu())
+    im_feat, cc_feat = torch.cat(im_feat), torch.cat(cc_feat)
+    win = 2 * vshift + 1
+    ccp = torch.nn.functional.pad(cc_feat, (0, 0, vshift, vshift))
+    dists = torch.stack([torch.nn.functional.pairwise_distance(im_feat[[i]].repeat(win, 1), ccp[i:i + win]) for i in range(n)], 1)
+    mean_d = dists.mean(1)
+    min_d, minidx = mean_d.min(0)
+    return int(vshift - minidx), float(min_d), float(mean_d.median() - min_d)
+
+
+def analyze(row, net, device):
+    vid = row["id"]
+    video_path = resolve_video_path(row)
+    vr = decord.VideoReader(str(video_path), num_threads=2)
+    fps = float(vr.get_avg_fps())
+    frames = vr.get_batch(list(range(len(vr)))).asnumpy()  # T,H,W,3 RGB
+    T, H, W = frames.shape[:3]
+
+    wav_path = DATA / "audio" / sub_dir(vid) / f"{vid}.wav"
+    wav_path.parent.mkdir(parents=True, exist_ok=True)
+    subprocess.run(["ffmpeg", "-loglevel", "error", "-nostdin", "-y", "-i", str(video_path), "-vn", "-ac", "1", "-ar", "16000",
+                    "-acodec", "pcm_s16le", str(wav_path)], check=True)
+    _, audio = wavfile.read(wav_path)
+
+    lm, nfaces = detect_landmarks(frames)
+    face_path = DATA / "face" / sub_dir(vid) / f"{vid}.npz"
+    face_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(face_path, lm=lm.astype(np.float16), nfaces=nfaces)
+
+    out = {"id": vid, "n_frames": T, "fps": fps, "height": H, "width": W, "dur_audio": len(audio) / 16000,
+           "face_missing": float((nfaces == 0).mean()), "multi_face": float((nfaces > 1).mean())}
+    valid = nfaces > 0
+    if valid.sum() < 0.5 * T:
+        return out
+    v = lm[valid]
+    face_h = v[:, LM_CHIN, 1] - v[:, LM_FOREHEAD, 1]
+    yaw = (v[:, LM_NOSE, 0] - v[:, LM_CHEEK_R, 0]) / (v[:, LM_CHEEK_L, 0] - v[:, LM_CHEEK_R, 0] + 1e-6)
+    mouth = (v[:, LM_LIP_LOW, 1] - v[:, LM_LIP_UP, 1]) / face_h
+    out.update({"face_h": float(np.median(face_h)), "yaw_ok": float(((yaw > 0.3) & (yaw < 0.7)).mean()),
+                "mouth_std": float(mouth.std())})
+
+    crops = syncnet_crops(frames, lm)
+    idx25 = (np.arange(int(T * 25 / fps)) * fps / 25).astype(int)
+    offset, min_dist, conf = syncnet_score(net, crops[idx25], audio, device)
+    out.update({"sync_offset": offset, "sync_dist": min_dist, "sync_conf": conf})
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--shard", type=int, default=0)
+    ap.add_argument("--num_shards", type=int, default=1)
+    ap.add_argument("--limit", type=int, default=None)
+    args = ap.parse_args()
+
+    rows = read_jsonl(MANIFEST / "s00_manifest.jsonl")
+    if args.limit:
+        rows = rows[:args.limit]
+    rows = shard(rows, args.shard, args.num_shards)
+    out_path = MANIFEST / "s01" / f"shard_{args.shard:04d}.jsonl"
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    done = {r["id"] for r in read_jsonl(out_path)} if out_path.exists() else set()
+    rows = [r for r in rows if r["id"] not in done]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    net = load_syncnet(device)
+    t0 = time.time()
+    with open(out_path, "a") as f:
+        for i, row in enumerate(rows):
+            try:
+                res = analyze(row, net, device)
+            except Exception as e:  # corrupt video etc.: record and continue
+                res = {"id": row["id"], "error": repr(e)[:300]}
+            f.write(json.dumps(res) + "\n")
+            f.flush()
+            if i % 50 == 0:
+                print(f"[shard {args.shard}] {i}/{len(rows)} {(time.time() - t0) / (i + 1):.2f}s/it", flush=True)
+
+
+if __name__ == "__main__":
+    main()

data/s02_make_clips.py

@@ -0,0 +1,198 @@
+"""Filter s01 results and cut training clips.
+
+Per clip (data/clips/<sub>/<clip_id>/):
+  video.mp4   512x512, 81 frames, 24fps (face-centred crop, fixed box per clip)
+  mask.png    512x512 uint8, 255 = region to generate (lower face, union over the clip)
+  ref.png     512x512 face frame from outside the clip window (same crop box)
+  audio.wav   16kHz mono, aligned to the clip (SyncNet offset corrected)
+Rows -> data/manifest/s02/shard_XXXX.jsonl
+"""
+import argparse
+import json
+import re
+import subprocess
+from pathlib import Path
+
+import cv2
+import decord
+import imageio_ffmpeg
+import numpy as np
+from scipy.io import wavfile
+
+from .common import DATA, MANIFEST, read_jsonl, resolve_video_path, shard
+
+NUM_FRAMES, SIZE, CROP_SCALE = 81, 512, 1.8
+LM_NOSE_BOTTOM, LM_CHEEK_R, LM_CHEEK_L, LM_CHIN, LM_LIP_UP, LM_LIP_LOW, LM_FOREHEAD = 2, 234, 454, 152, 13, 14, 10
+FILTERS = {  # decided from the s01 distribution, see STATE.md
+    "face_missing_max": 0.02, "multi_face_max": 0.02, "face_h_min": 150, "yaw_ok_min": 0.7,
+    "sync_conf_min": 6.0, "sync_offset_range": (-4, 2), "mouth_std_min": 0.01,
+}
+
+
+def keep(r):
+    f = FILTERS
+    return ("error" not in r and "sync_conf" in r and r["face_missing"] <= f["face_missing_max"]
+            and r["multi_face"] <= f["multi_face_max"] and r["face_h"] >= f["face_h_min"] and r["yaw_ok"] >= f["yaw_ok_min"]
+            and r["sync_conf"] >= f["sync_conf_min"] and f["sync_offset_range"][0] <= r["sync_offset"] <= f["sync_offset_range"][1]
+            and r["mouth_std"] >= f["mouth_std_min"] and r["n_frames"] >= NUM_FRAMES)
+
+
+def clean_prompt(prompt, speech_text):
+    """Drop sentences that carry the spoken line (it would conflict with new audio at dubbing time)."""
+    prompt = prompt.replace(speech_text, "")
+    parts = re.split(r"(?<=[。！？.!?])\s*", prompt)
+    bad = re.compile(r"台词|口型|说道|说：|lip|speak|says|saying|dialogue|words", re.I)
+    return " ".join(p for p in parts if p.strip() and not bad.search(p)).strip()
+
+
+def crop_box(lm, H, W):
+    cx = (np.nanmin(lm[:, :, 0]) + np.nanmax(lm[:, :, 0])) / 2
+    cy = (np.nanmin(lm[:, :, 1]) + np.nanmax(lm[:, :, 1])) / 2
+    face = np.nanmedian(np.maximum(np.ptp(lm[:, :, 0], 1), np.ptp(lm[:, :, 1], 1)))
+    side = int(min(face * CROP_SCALE, H, W))
+    x0 = int(np.clip(cx - side / 2, 0, W - side))
+    y0 = int(np.clip(cy - side / 2, 0, H - side))
+    return x0, y0, side
+
+
+def lower_face_mask(lm, box):
+    x0, y0, side = box
+    s = SIZE / side
+    face_h = np.nanmedian(lm[:, LM_CHIN, 1] - lm[:, LM_FOREHEAD, 1])
+    top = np.nanmin(lm[:, LM_NOSE_BOTTOM, 1])
+    bottom = np.nanmax(lm[:, LM_CHIN, 1]) + 0.1 * face_h
+    left, right = np.nanmin(lm[:, LM_CHEEK_R, 0]), np.nanmax(lm[:, LM_CHEEK_L, 0])
+    mask = np.zeros((SIZE, SIZE), np.uint8)
+    pts = np.array([[left, top], [right, bottom]])
+    (l, t), (r, b) = ((pts - [x0, y0]) * s).round().astype(int)
+    mask[max(t, 0):min(b, SIZE), max(l, 0):min(r, SIZE)] = 255
+    return mask
+
+
+def write_video(path, frames):
+    proc = subprocess.Popen([imageio_ffmpeg.get_ffmpeg_exe(), "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{SIZE}x{SIZE}",
+                             "-r", "24", "-i", "-", "-c:v", "libx264", "-crf", "15", "-pix_fmt", "yuv420p", str(path)], stdin=subprocess.PIPE)
+    proc.stdin.write(np.ascontiguousarray(frames).tobytes())
+    proc.stdin.close()
+    assert proc.wait() == 0
+
+
+def load_selection(path):
+    """Return source_id -> selected clip_ids from JSON, JSONL, or text."""
+    path = Path(path)
+    if path.suffix == ".json":
+        value = json.loads(path.read_text())
+        values = value.keys() if isinstance(value, dict) else value
+    else:
+        values = []
+        with path.open() as handle:
+            for line in handle:
+                if not line.strip():
+                    continue
+                values.append(json.loads(line) if line.lstrip().startswith("{") else line.strip())
+    selected = {}
+    for value in values:
+        clip_id = value["clip_id"] if isinstance(value, dict) else str(value)
+        source_id = value.get("id") if isinstance(value, dict) else None
+        source_id = source_id or clip_id.rsplit("_", 1)[0]
+        selected.setdefault(source_id, set()).add(clip_id)
+    if not selected:
+        raise ValueError(f"Empty selection: {path}")
+    return selected
+
+
+def make_clips(row, s01, selected=None, done=None):
+    vid = row["id"]
+    sub = vid[-6:-3]
+    face = np.load(DATA / "face" / sub / f"{vid}.npz")
+    lm = face["lm"].astype(np.float32)
+    T = len(lm)
+    vr = decord.VideoReader(str(resolve_video_path(row)), num_threads=2)
+    frames = vr.get_batch(list(range(T))).asnumpy()
+    H, W = frames.shape[1:3]
+    _, audio = wavfile.read(DATA / "audio" / sub / f"{vid}.wav")
+    mouth = (lm[:, LM_LIP_LOW, 1] - lm[:, LM_LIP_UP, 1]) / (lm[:, LM_CHIN, 1] - lm[:, LM_FOREHEAD, 1])
+    mouth = np.nan_to_num(mouth, nan=-1)
+    # best-matching audio for video time t is at t - offset/25 s (offset < 0: audio is late)
+    audio_shift = -s01["sync_offset"] / 25.0
+    starts = [0] if T - NUM_FRAMES < 16 else [0, T - NUM_FRAMES]
+    out = []
+    for start in starts:
+        clip_id = f"{vid}_{start:03d}"
+        if selected is not None and clip_id not in selected:
+            continue
+        if done is not None and clip_id in done:
+            continue
+        win = slice(start, start + NUM_FRAMES)
+        box = crop_box(lm[win], H, W)
+        x0, y0, side = box
+        outside = np.r_[0:start, start + NUM_FRAMES:T]
+        ref_idx = int(outside[np.argmax(mouth[outside])]) if len(outside) else (T - 1 if start == 0 else 0)
+        crop = lambda f: cv2.resize(f[y0:y0 + side, x0:x0 + side], (SIZE, SIZE), interpolation=cv2.INTER_AREA)
+
+        d = DATA / "clips" / sub / clip_id
+        d.mkdir(parents=True, exist_ok=True)
+        write_video(d / "video.mp4", np.stack([crop(f) for f in frames[win]]))
+        cv2.imwrite(str(d / "mask.png"), lower_face_mask(lm[win], box))
+        cv2.imwrite(str(d / "ref.png"), cv2.cvtColor(crop(frames[ref_idx]), cv2.COLOR_RGB2BGR))
+        a0 = int(round((start / 24.0 + audio_shift) * 16000))
+        n = int(round(NUM_FRAMES / 24.0 * 16000))
+        seg = np.zeros(n, np.int16)
+        src = audio[max(a0, 0):max(a0 + n, 0)]
+        seg[max(-a0, 0):max(-a0, 0) + len(src)] = src[:n - max(-a0, 0)]
+        wavfile.write(d / "audio.wav", 16000, seg)
+        out.append({"clip_id": clip_id, "id": vid, "dir": str(d.relative_to(DATA / "clips")), "start": start, "ref_idx": ref_idx,
+                    "crop_box": box, "sync_conf": s01["sync_conf"], "sync_offset": s01["sync_offset"],
+                    "lang": row.get("speech_language", row.get("language", "")),
+                    "prompt": clean_prompt(row.get("prompt", ""), row.get("speech_text", ""))})
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--shard", type=int, default=0)
+    ap.add_argument("--num_shards", type=int, default=1)
+    ap.add_argument("--limit", type=int, default=None)
+    ap.add_argument(
+        "--selection",
+        help="Optional JSON/JSONL/text clip selection; unlisted windows are not materialized",
+    )
+    args = ap.parse_args()
+
+    selection = load_selection(args.selection) if args.selection else None
+    s01 = {}
+    for p in sorted((MANIFEST / "s01").glob("shard_*.jsonl")):
+        s01.update({r["id"]: r for r in read_jsonl(p)})
+    rows = [r for r in read_jsonl(MANIFEST / "s00_manifest.jsonl") if r["id"] in s01 and keep(s01[r["id"]])]
+    if selection is not None:
+        rows = [r for r in rows if r["id"] in selection]
+    if args.limit:
+        rows = rows[:args.limit]
+    rows = shard(rows, args.shard, args.num_shards)
+    out_path = MANIFEST / "s02" / f"shard_{args.shard:04d}.jsonl"
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    done = {
+        r["clip_id"] for r in read_jsonl(out_path) if "clip_id" in r
+    } if out_path.exists() else set()
+    with open(out_path, "a") as f:
+        for i, row in enumerate(rows):
+            try:
+                clips = make_clips(
+                    row,
+                    s01[row["id"]],
+                    selected=selection.get(row["id"]) if selection is not None else None,
+                    done=done,
+                )
+            except Exception as e:
+                clips = [{"id": row["id"], "error": repr(e)[:300]}]
+            for c in clips:
+                f.write(json.dumps(c, ensure_ascii=False) + "\n")
+                if "clip_id" in c:
+                    done.add(c["clip_id"])
+            f.flush()
+            if i % 50 == 0:
+                print(f"[shard {args.shard}] {i}/{len(rows)}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

examples/man_laughing_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f53d0999a39251346563a4de62384920c28a67c053ae91fe05c948bf04644c5e
+size 144078

examples/man_laughing_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a09134a1e470c253e9241521887c7bd01714eaf45a8a16159b6bd325ec8c0eaa
+size 1003686

examples/man_smiling_studio_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da0044fcc67fd95ba6b3e75c8f5507905c046100b8074fb13cb3da82ac6f6c05
+size 144078

examples/man_smiling_studio_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a54a7fe60ceb0be72dc6cd2f6dd3b91ed3b83db919a0e618cd91168e011f9cb
+size 776166

examples/woman_laughing_studio_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99cb56c5ccb274afeed195232dcd65745ed4eab9559a4b0f34670e32defada7b
+size 144078

examples/woman_laughing_studio_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:911cf4085ab10fd231136a1cb7c72e4e2e0cc45da5e8a0672a49924646c26b0c
+size 716276

requirements.txt

@@ -0,0 +1,18 @@
+numpy==1.26.4
+mediapipe==0.10.21
+opencv-contrib-python-headless==4.11.0.86
+decord
+librosa
+soundfile
+scipy
+python-speech-features
+einops
+ftfy
+regex
+safetensors
+imageio-ffmpeg
+Pillow
+tqdm
+accelerate==1.10.1
+transformers==4.56.2
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt212-cu130-cp310/flash_attn-2.8.3-cp310-cp310-linux_x86_64.whl

viddub/__init__.py

@@ -0,0 +1,3 @@
+"""VidDub model, training, and inference implementation."""
+
+__version__ = "1.0.0"

viddub/causal/__init__.py

@@ -0,0 +1 @@
+"""Checkpoint-compatible, block-causal Wan-Dub training and sampling."""

viddub/causal/art_alignment.py

@@ -0,0 +1,44 @@
+"""Trailing-window Art-REPA for a causal prefix; no future hidden states.
+
+This module is prepared for second-wave experiments. It is not enabled in A0-A3.
+The last four video latents map to the last sixteen decoded frames once at
+least five video latents are available. Reference positions are excluded.
+"""
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class CausalArtAlignment(nn.Module):
+    def __init__(self,dim=1536,embedding_dim=2048):
+        super().__init__()
+        self.dim=dim
+        self.proj=nn.Sequential(nn.Linear(4*dim,embedding_dim),nn.SiLU(),
+                                nn.Linear(embedding_dim,embedding_dim),nn.SiLU(),
+                                nn.Linear(embedding_dim,embedding_dim))
+
+    def forward(self,hidden,pixel_mask,target,*,num_ref=1,history_end=0,warmup=False):
+        """hidden=[B,F,Htokens,Wtokens,C], target=frozen GT visual embedding.
+
+        All hidden positions through history_end are detached, as they would
+        be for a generated-history cache. warmup detaches every hidden state.
+        Caller must pass a prefix ending at the supervised block boundary.
+        """
+        b,frames,h,w,c=hidden.shape
+        if c!=self.dim or target.shape!=(b,self.proj[-1].out_features):
+            raise ValueError('Art hidden/target dimensions differ')
+        if frames-num_ref<5:
+            # Keep trainable projector represented for DDP without inventing
+            # a 16-frame target at the first short block.
+            return sum(parameter.sum()*0 for parameter in self.proj.parameters())
+        token_mask=F.adaptive_max_pool2d(pixel_mask.float()[None,None],(h,w))[0,0]>0
+        if not token_mask.any():raise ValueError('Empty Art mouth mask')
+        values=hidden[:,frames-4:]
+        if warmup:values=values.detach()
+        elif history_end:
+            cut=max(0,min(4,history_end-(frames-4)))
+            values=torch.cat([values[:,:cut].detach(),values[:,cut:]],1)
+        pooled=values[:,:,token_mask].mean(2).flatten(1)
+        pred=F.normalize(self.proj(pooled).float(),dim=-1)
+        target=F.normalize(target.detach().float(),dim=-1)
+        return (1-(pred*target).sum(-1)).mean()

viddub/causal/cf/NOTICE

@@ -0,0 +1,18 @@
+Wan-Dub Causal Forcing++ adaptation
+
+The flow-matching scheduler compatibility code is adapted from the official
+Causal-Forcing repository, commit da3ddf1a590f30c9edc568e11c2f77170995c31c,
+under Apache-2.0. The license text is retained under
+dependencies/licenses/Apache-2.0.txt.
+
+The adjacent-timestep consistency algorithm and re-noising sampler in this
+directory are adaptations of model/naive_consistency.py,
+utils/wan_wrapper.py, pipeline/causal_inference.py and
+pipeline/self_forcing_training.py from that repository. The DMD adapter follows
+model/dmd.py with documented differences in gradient truncation and conditioning.
+
+Changes by the Wan-Dub integration: VACE/audio conditions, independent reference,
+split-backward CD, sampled-chunk DMD gradients, explicit sampler metadata,
+finite jobs, full training-state restore, task-specific GT/Sync objectives and
+source-matched validation. See the repository README and
+`docs/ARCHITECTURE.md` for the public implementation scope.

viddub/causal/cf/__init__.py

@@ -0,0 +1 @@
+"""Causal Forcing++ adaptation for the existing Wan-Dub model and conditions."""

viddub/causal/cf/consistency.py

@@ -0,0 +1,123 @@
+"""Adapted from official model/naive_consistency.py (Apache-2.0).
+
+Changes: Wan-Dub conditions/reference; optional frozen history caches; chunkwise
+backward to bound memory. All seven losses can be accumulated with full student
+history gradients, matching the upstream teacher-forcing attention dependency.
+"""
+from contextlib import contextmanager
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .upstream import FlowMatchScheduler
+from ..model import HistoryCache
+from ..sampling import commit, evaluating, predict_chunk
+
+
+def flow_to_x0(xt, velocity, timestep, scheduler):
+    """Same nearest-grid and float64 arithmetic as official WanDiffusionWrapper."""
+    grid = scheduler.timesteps.to(device=xt.device, dtype=torch.float64)
+    index = (grid - torch.as_tensor(timestep, device=xt.device).double()).abs().argmin()
+    sigma = scheduler.sigmas.to(device=xt.device, dtype=torch.float64)[index]
+    return (xt.double() - sigma * velocity.double()).to(velocity.dtype)
+
+
+class ModelEMA:
+    """Frozen target model, updated every optimizer step and included in resume."""
+    def __init__(self, model, decay):
+        self.model = model.eval().requires_grad_(False)
+        self.decay, self.updates = decay, 0
+
+    @torch.no_grad()
+    def update(self, student):
+        for (_, target), (_, source) in zip(self.model.named_parameters(), student.named_parameters()):
+            target.mul_(self.decay).add_(source.detach().to(target), alpha=1-self.decay)
+        for (_, target), (_, source) in zip(self.model.named_buffers(), student.named_buffers()):
+            target.copy_(source)
+        self.updates += 1
+
+    def state_dict(self):
+        return {'model': self.model.state_dict(), 'decay': self.decay, 'updates': self.updates}
+
+    def load_state_dict(self, state):
+        if state['decay'] != self.decay:
+            raise ValueError('EMA decay differs on resume')
+        self.model.load_state_dict(state['model'], strict=True)
+        self.updates = state['updates']
+
+
+@torch.no_grad()
+def clean_history_cache(model, clean, conditions, end):
+    cache = HistoryCache()
+    with evaluating(model):
+        for a, b in model.config.blocks(clean.shape[2]):
+            if b > end:
+                break
+            commit(model, clean[:, :, a:b], conditions, a, cache)
+    if cache.frames != end:
+        raise ValueError('History endpoint must be a block boundary')
+    return cache
+
+
+class CausalCD(nn.Module):
+    def __init__(self, model, teacher, ema, cfg, quality=None):
+        super().__init__()
+        self.model = model
+        # Do not register frozen targets in DDP or the student optimizer.
+        self.targets = [teacher, ema]
+        self.quality = [quality]
+        self.cfg = cfg
+        self.grid = FlowMatchScheduler(shift=cfg['shift'], sigma_min=0., extra_one_step=True)
+        self.grid.set_timesteps(cfg['cd_grid_steps'])
+        self.conversion_grid = FlowMatchScheduler(shift=cfg['shift'], sigma_min=0., extra_one_step=True)
+        self.conversion_grid.set_timesteps(1000, training=True)
+
+    def sample_input(self, clean, index, noise):
+        if not 0 <= index < self.cfg['cd_grid_steps'] - 1:
+            raise ValueError('CD adjacent index outside nonzero upstream grid')
+        # Official add_noise uses the exact 48-point sigma; the transformer
+        # timesteps are cast to the configured compute dtype (BF16 upstream).
+        s = self.grid.sigmas[index].to(device=clean.device, dtype=torch.float32)
+        xt = ((1-s.float()) * clean.float() + s.float()*noise.float()).to(clean.dtype)
+        t = self.grid.timesteps[index].to(clean)
+        tn = self.grid.timesteps[index+1].to(clean)
+        return xt, t, tn
+
+    def forward(self, clean, c, *, xt, timestep, next_timestep, chunk_index):
+        a, b = self.model.config.video_blocks(clean.shape[2])[chunk_index]
+        current = xt[:, :, a:b]
+        prefix = clean[:, :, :a].detach()
+        teacher, ema = self.targets
+        with torch.no_grad():
+            teacher_cache = clean_history_cache(teacher, clean, c, a) if self.cfg['cd_frozen_cache'] else None
+            velocity = predict_chunk(teacher, current, float(timestep)/1000, prefix, c, a, teacher_cache)
+            if self.cfg['teacher_cfg'] != 1:
+                null_c = dict(c, drop_audio=torch.ones(clean.shape[0],device=clean.device,dtype=torch.bool))
+                # Conditional and unconditional histories must be different.
+                velocity_null = predict_chunk(teacher, current, float(timestep)/1000, prefix, c, a,
+                                              drop_audio=null_c['drop_audio'])
+                velocity = velocity_null + self.cfg['teacher_cfg'] * (velocity-velocity_null)
+            dt = (timestep-next_timestep)/1000
+            next_xt = current - dt * velocity
+            del teacher_cache
+            ema_cache = clean_history_cache(ema, clean, c, a) if self.cfg['cd_frozen_cache'] else None
+            ema_v = predict_chunk(ema, next_xt, float(next_timestep)/1000, prefix, c, a, ema_cache)
+            target = flow_to_x0(next_xt, ema_v, next_timestep, self.conversion_grid).detach()
+            del ema_cache
+        # No student cache by default: preserve parameter gradients through
+        # clean-history forward, as in the official doubled-sequence TF mask.
+        cache = clean_history_cache(self.model, clean, c, a) if self.cfg['cd_student_cache'] else None
+        velocity = predict_chunk(self.model, current, float(timestep)/1000, prefix, c, a, cache,
+                                 gradient_checkpointing=self.cfg['gradient_checkpointing'])
+        estimate = flow_to_x0(current, velocity, timestep, self.conversion_grid)
+        cd_loss = F.mse_loss(estimate.float(), target.float())
+        loss = cd_loss
+        stats = {'cd': cd_loss.detach(), 'sigma': float(timestep)/1000,
+                 'chunk': float(chunk_index), 'aux_loss': 0.}
+        if self.quality[0] is not None:
+            available = torch.cat([prefix, estimate], 2)
+            aux, aux_stats = self.quality[0](available, c, a, b, float(timestep)/1000)
+            loss = loss + self.cfg['cd_aux_scale'] * aux
+            stats.update({k: v.detach() for k,v in aux_stats.items()})
+            stats['aux_loss'] = aux.detach()
+        return loss, stats

viddub/causal/cf/dmd.py

@@ -0,0 +1,98 @@
+"""Asymmetric DMD aligned to pinned CF sampling and score-query conventions.
+
+The real/fake models remain bidirectional. A single exit step is used for all
+chunks and both fake/generator rollouts use the same exit-step distribution.
+Wan-Dub differences: audio CFG (standard null + cfg*(cond-null) convention),
+optional mouth weighting/quality/ODE retention, extra clean reference excluded
+from video loss. No pretrained T2V checkpoint is substituted for Wan-Dub.
+"""
+import torch
+import torch.distributed as dist
+from .sampling import cf_rollout
+from .upstream import FlowMatchScheduler
+from ..training.common import flow_noise, spatial_weight, weighted_mse
+from ..training.dmd import DMDObjective
+
+
+def shared_exit(steps, device):
+    value=torch.randint(steps,(1,),device=device)
+    if dist.is_available() and dist.is_initialized():dist.broadcast(value,src=0)
+    return int(value)
+
+
+def score_sigma(device, shift=5., minimum=.02, maximum=.98):
+    # Official DMD samples integer indices, shifts, clamps, then add_noise
+    # maps the timestep to the nearest point on the 1000-point scheduler.
+    t=torch.randint(0,1000,(),device=device).float()/1000
+    t=(shift*t/(1+(shift-1)*t)).clamp(minimum,maximum)*1000
+    grid=FlowMatchScheduler(shift=shift,sigma_min=0.,extra_one_step=True)
+    grid.set_timesteps(1000,training=True)
+    index=(grid.timesteps.to(device)-t).abs().argmin()
+    return grid.sigmas.to(device)[index]
+
+
+class CFAsymmetricDMD(DMDObjective):
+    def dmd(self,x0,c,chunk_index=None,dmd_mode='combined'):
+        if dmd_mode != 'combined':
+            raise ValueError('CF asymmetric DMD does not implement split-quality backward')
+        teacher,fake,anchor=self.scores
+        if teacher is None or fake is None or teacher.causal or fake.causal:
+            raise ValueError('CF asymmetric DMD requires bidirectional real and fake scores')
+        ranges=self.model.config.video_blocks(x0.shape[2])
+        i=int(torch.randint(len(ranges),())) if chunk_index is None else chunk_index
+        a,b=ranges[i]
+        exit_step=(self.student_steps-1 if self.dmd_rollout_mode=='final_step'
+                   else shared_exit(self.student_steps,x0.device))
+        seed=int(torch.randint(2**31-1,(),device=x0.device))
+        generator=torch.Generator(device=x0.device).manual_seed(seed)
+        # Sample one chunk for parameter gradients. This is a documented
+        # memory-bounded unbiased chunk estimator, not full upstream backprop.
+        gen=cf_rollout(self.model,x0,c,steps=self.student_steps,shift=self.shift,
+            generator=generator,use_cache=self.use_cache,grad_chunk=i,grad_exit=exit_step,
+            gradient_checkpointing=self.gradient_checkpointing)
+        with torch.no_grad():
+            s=score_sigma(x0.device,self.shift,self.matching_min,self.matching_max)
+            noisy=flow_noise(gen.detach(),s,torch.randn_like(gen))
+            ts=noisy.new_full((noisy.shape[0],noisy.shape[2]),float(s)*1000)
+            real_v=self._score(teacher,noisy,ts,c)
+            if self.teacher_cfg!=1:
+                null=torch.ones(noisy.shape[0],device=x0.device,dtype=torch.bool)
+                null_v=self._score(teacher,noisy,ts,c,null)
+                real_v=null_v+self.teacher_cfg*(real_v-null_v)
+            real_x=noisy.float()-s*real_v.float()
+            fake_x=noisy.float()-s*self._score(fake,noisy,ts,c).float()
+            n=self.model.config.num_ref
+            denom=(gen[:,:,n:].float()-real_x[:,:,n:]).abs().mean((1,2,3,4),keepdim=True).clamp_min(1e-6)
+            direction=(fake_x-real_x)/denom
+            if not torch.isfinite(direction).all():raise FloatingPointError('Nonfinite DMD direction')
+            target=gen.detach().float()-direction
+        weight=spatial_weight(c['vace_context'][:,:,a:b],self.mouth_weight)
+        loss=.5*weighted_mse(gen[:,:,a:b],target[:,:,a:b],weight)
+        stats={'dmd':loss.detach(),'chunk':float(i),'exit_step':float(exit_step),
+               'score_sigma':float(s),'direction_rms':direction[:,:,n:].square().mean().sqrt()}
+        if self.dmd_retention_weight:
+            if anchor is None:raise ValueError('Missing frozen anchor')
+            if self.dmd_score_offload:anchor.to(x0.device)
+            with torch.no_grad():
+                fixed=cf_rollout(anchor,x0,c,steps=self.student_steps,shift=self.shift,
+                    generator=torch.Generator(device=x0.device).manual_seed(seed),
+                    use_cache=self.use_cache,grad_exit=exit_step)
+            if self.dmd_score_offload:anchor.cpu();torch.cuda.empty_cache()
+            retain=weighted_mse(gen[:,:,a:b],fixed[:,:,a:b],weight)
+            loss+=self.dmd_retention_weight*retain;stats['retention']=retain.detach()
+        if self.quality[0] is not None and (not c.get('cross',False) or self.rollout_quality_cross):
+            aux,aux_stats=self.quality[0](gen[:,:,:b],c,a,b,0.)
+            loss+=self.rollout_quality_weight*aux
+            stats.update({f'rollout_{k}':v.detach() for k,v in aux_stats.items()})
+        return loss,stats
+
+    def fake(self,generated,c):
+        if generated is None:raise ValueError('Fake score requires generated samples')
+        gen=generated.detach();s=score_sigma(gen.device,self.shift,self.matching_min,self.matching_max)
+        noise=torch.randn_like(gen);noisy=flow_noise(gen,s,noise)
+        ts=gen.new_full((gen.shape[0],gen.shape[2]),float(s)*1000)
+        pred=self._score(self.model,noisy,ts,c,grad=True)
+        n=self.model.config.num_ref
+        loss=weighted_mse(pred[:,:,n:],(noise-gen)[:,:,n:],
+                           spatial_weight(c['vace_context'][:,:,n:],self.mouth_weight))
+        return loss,{'fake_flow':loss.detach(),'score_sigma':float(s)}

viddub/causal/cf/sampling.py

@@ -0,0 +1,54 @@
+"""Official CF few-step rule: x0 prediction then independent re-noising.
+
+Wan-Dub implementation of pipeline/causal_inference.py from the pinned source.
+Unlike Euler sampling, this has independent Gaussian noise at each transition.
+"""
+import torch
+from .upstream import FlowMatchScheduler
+from .consistency import flow_to_x0
+from ..model import HistoryCache
+from ..sampling import evaluating, commit, predict_chunk
+
+
+def cf_rollout(model, shape_reference, c, *, steps=4, shift=5., noise=None,
+               generator=None, use_cache=True, grad_chunk=None, grad_exit=None,
+               gradient_checkpointing=False):
+    if steps < 1 or 1000 % steps:
+        raise ValueError('CF schedule requires steps to divide 1000')
+    scheduler = FlowMatchScheduler(shift=shift, sigma_min=0., extra_one_step=True)
+    scheduler.set_timesteps(1000, training=True)
+    timesteps = scheduler.timesteps[torch.arange(0,1000,1000//steps)].to(shape_reference.device)
+    noise = torch.randn(shape_reference.shape,device=shape_reference.device,dtype=shape_reference.dtype,
+                        generator=generator) if noise is None else noise
+    if noise.shape != shape_reference.shape:
+        raise ValueError('Noise shape differs')
+    cache = HistoryCache() if use_cache else None
+    prefix = shape_reference[:, :, :model.config.num_ref].detach()
+    output = [prefix]
+    if grad_exit is not None and not 0 <= grad_exit < steps:
+        raise ValueError('Invalid exit step')
+    with evaluating(model):
+        if prefix.shape[2]:
+            commit(model, prefix, c, 0, cache)
+        for index, (a,b) in enumerate(model.config.video_blocks(shape_reference.shape[2])):
+            current = noise[:, :, a:b]
+            final = steps-1 if grad_exit is None else grad_exit
+            for j in range(final+1):
+                keep = index == grad_chunk and j == final
+                with torch.set_grad_enabled(keep):
+                    # Upstream quantizes the warped timestep to integer at inference.
+                    t = timesteps[j].long()
+                    v = predict_chunk(model,current,float(t)/1000,prefix,c,a,cache,
+                                      gradient_checkpointing=gradient_checkpointing,detach_history=True)
+                    estimate = flow_to_x0(current,v,t,scheduler)
+                if j < final:
+                    with torch.no_grad():
+                        tnext = timesteps[j+1].long()
+                        step_index = (scheduler.timesteps.to(tnext.device)-tnext).abs().argmin()
+                        sigma = scheduler.sigmas[step_index].to(estimate)
+                        eps = torch.randn(estimate.shape,device=estimate.device,dtype=estimate.dtype,generator=generator)
+                        current = ((1-sigma)*estimate.detach()+sigma*eps).to(shape_reference.dtype)
+            output.append(estimate)
+            commit(model,estimate,c,a,cache)
+            prefix = torch.cat([prefix.detach(),estimate.detach()],2)
+    return torch.cat(output,2)

viddub/causal/cf/train.py

@@ -0,0 +1,168 @@
+"""Finite, resumable CF++ causal-CD trainer with Wan-Dub conditioning.
+
+Uses the official online adjacent teacher-step objective, AdamW and EMA.
+All blocks contribute per clip; individual backward calls bound activation
+memory without detaching student clean-history parameter paths by default.
+"""
+import argparse
+import copy
+import datetime
+import hashlib
+import json
+import random
+import time
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+
+from .consistency import CausalCD, ModelEMA
+from .upstream import provenance as upstream_provenance
+from ..runtime import (atomic_json, autocast, checkpoint_paths, export_weights, fingerprint,
+                       first_item, load_components, local_imports, make_dataset, prepare_clip,
+                       read_config, require_causal_metadata)
+from ..train import check_loss, finite_update, optimizer
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--config',required=True)
+    parser.add_argument('--init-dir',required=True)
+    parser.add_argument('--init-step',type=int,required=True)
+    parser.add_argument('--output',required=True)
+    parser.add_argument('--resume')
+    parser.add_argument('--stop-after-step',type=int)
+    parser.add_argument('--check-only',action='store_true')
+    args = parser.parse_args()
+    cfg = read_config(args.config)
+    if cfg['stage'] != 'cd' or cfg['grad_accum'] != 1 or cfg['p_cross'] != 0 or cfg['art_weight']:
+        raise ValueError('CD pilot requires stage=cd, grad_accum=1, same data, Art disabled')
+    if cfg['warmup'] or cfg['lr_schedule'] != 'constant':
+        raise ValueError('This CF CD trainer uses the official constant LR without warmup')
+    final = cfg['max_steps'] if args.stop_after_step is None else args.stop_after_step
+    if not 0 < final <= cfg['max_steps']:
+        raise ValueError('Invalid operational stop')
+    out = Path(args.output)
+    if not args.resume and out.exists() and any(out.iterdir()):
+        raise FileExistsError(f'Output is not empty: {out}')
+    paths = checkpoint_paths(args.init_dir,args.init_step)
+    require_causal_metadata(args.init_dir,args.init_step,cfg,('ar',))
+    dataset, rows = make_dataset(cfg)
+    record = {'student_init': fingerprint(paths), 'upstream':upstream_provenance(),
+              'rows_sha256':hashlib.sha256(json.dumps(rows,sort_keys=True).encode()).hexdigest(),
+              'adapter_sources':fingerprint({n:Path(__file__).parent/n for n in
+                                           ['train.py','consistency.py','sampling.py','upstream.py']})}
+    resume = None
+    if args.resume:
+        if Path(args.resume).resolve().parent != out.resolve():
+            raise ValueError('Resume state must belong to this run')
+        resume = json.loads((Path(args.resume)/'complete.json').read_text())
+        if read_config(resume['config']) != cfg or resume['provenance'] != record:
+            raise ValueError('Resume config, code or provenance differs')
+    if args.check_only:
+        print(json.dumps({'status':'validated_no_gpu','samples':len(dataset),'provenance':record},indent=2))
+        return
+    local_imports()
+    from accelerate import Accelerator, DistributedDataParallelKwargs, InitProcessGroupKwargs
+    from viddub.full_context.runtime import ZeroLatents
+    accel = Accelerator(mixed_precision='no',gradient_accumulation_steps=1,
+        kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=False,gradient_as_bucket_view=True),
+                        InitProcessGroupKwargs(timeout=datetime.timedelta(minutes=20))])
+    if accel.device.type != 'cuda':
+        raise RuntimeError('Training requires CUDA')
+    random.seed(cfg['seed']+accel.process_index);torch.manual_seed(cfg['seed']+accel.process_index)
+    pipe, whisper, student = load_components(paths,cfg,accel.device)
+    student.requires_grad_(True).eval()
+    teacher = copy.deepcopy(student).requires_grad_(False).eval()
+    ema = ModelEMA(copy.deepcopy(student),cfg['cd_ema_decay'])
+    quality = None
+    if any(cfg[k] for k in ['pixel_weight','lpips_weight','keep_weight','sync_weight']):
+        from ..quality import GTQuality
+        quality = GTQuality(pipe,cfg,accel.device)
+    obj = CausalCD(student,teacher,ema.model,cfg,quality)
+    opt = optimizer(student,cfg)
+    loader = torch.utils.data.DataLoader(dataset,batch_size=1,shuffle=True,collate_fn=first_item,
+        num_workers=cfg['workers'],generator=torch.Generator().manual_seed(cfg['seed']))
+    obj,opt,loader = accel.prepare(obj,opt,loader)
+    accel.register_for_checkpointing(ema)
+    step=epoch=consumed=0
+    if resume:
+        if resume['world_size'] != accel.num_processes:
+            raise ValueError('Resume requires same world size')
+        accel.load_state(args.resume)
+        step,epoch,consumed=(resume[k] for k in ('step','epoch','consumed_batches'))
+    else:
+        random.seed(cfg['seed']+accel.process_index);torch.manual_seed(cfg['seed']+accel.process_index)
+    if accel.is_main_process:
+        out.mkdir(parents=True,exist_ok=True)
+        atomic_json(out/'run.json',{'config':cfg,'provenance':record,'world_size':accel.num_processes,
+                                   'scope':'diagnostic_CD_pilot_not_AR_promotion'})
+    accel.wait_for_everyone()
+    zero = ZeroLatents(pipe)
+    torch.cuda.reset_peak_memory_stats()
+    started=time.monotonic()
+    def save():
+        state=out/f'state_step{step:06d}'
+        exists=accel.reduce(torch.tensor(int(state.exists()),device=accel.device),'sum')
+        if int(exists):raise FileExistsError(state)
+        full = step == cfg['max_steps'] or step == final or step % cfg['save_state_every'] == 0
+        if full:accel.save_state(str(state))
+        if accel.is_main_process:
+            export_weights(student,out,step,cfg,record)
+            export_weights(ema.model,out/'ema',step,cfg,{**record,'export':'ema','ema_updates':ema.updates})
+            if full:
+                atomic_json(state/'complete.json',{'step':step,'epoch':epoch,'consumed_batches':consumed,
+                    'config':cfg,'provenance':record,'world_size':accel.num_processes,'ema_updates':ema.updates})
+        accel.wait_for_everyone()
+    while step < final:
+        if hasattr(loader,'set_epoch'):loader.set_epoch(epoch)
+        if getattr(loader,'generator',None) is not None:loader.generator.manual_seed(cfg['seed']+epoch)
+        active=accel.skip_first_batches(loader,consumed) if consumed else loader
+        for item in active:
+            consumed += 1
+            clean,c = prepare_clip(item,pipe,whisper,zero,cfg)
+            c['cross']=False
+            # Identical RNG consumption in pure CD and GT auxiliary variants.
+            index=random.randrange(cfg['cd_grid_steps']-1)
+            noise=torch.randn_like(clean)
+            raw=accel.unwrap_model(obj)
+            xt,t,tn=raw.sample_input(clean,index,noise)
+            ranges=student.config.video_blocks(clean.shape[2])
+            indices=list(range(len(ranges))) if cfg['cd_all_chunks'] else [int(torch.randint(len(ranges),()))]
+            stats={};value=clean.new_zeros((),dtype=torch.float32)
+            opt.zero_grad(set_to_none=True)
+            for j,i in enumerate(indices):
+                # Backward per chunk; accumulate parameter gradients before one
+                # DDP all-reduce and optimizer step for the complete clip.
+                context=accel.no_sync(obj) if j<len(indices)-1 else nullcontext()
+                with context:
+                    with autocast(accel.device,cfg['precision']):
+                        loss,metrics=obj(clean,c,xt=xt,timestep=t,next_timestep=tn,chunk_index=i)
+                        loss=loss/len(indices)
+                    check_loss(accel,loss)
+                    accel.backward(loss)
+                value += loss.detach()
+                for key,v in metrics.items():stats[key]=stats.get(key,0.)+float(v)/len(indices)
+                del loss,metrics
+            norm=finite_update(accel,obj,opt,cfg)
+            ema.update(student)
+            step += 1
+            if step%cfg['log_every']==0 or step==final:
+                keys=sorted(stats)
+                values=torch.tensor([stats[k] for k in keys],device=accel.device)
+                values=accel.reduce(values,'mean').tolist()
+                log={'step':step,'loss':float(accel.reduce(value,'mean')),'grad_norm':norm,
+                    'elapsed_seconds':time.monotonic()-started,'ema_updates':ema.updates,
+                    'blocks_per_clip':len(indices),'grid_index':index,
+                    'peak_allocated_gib':torch.cuda.max_memory_allocated()/1024**3,
+                    'peak_reserved_gib':torch.cuda.max_memory_reserved()/1024**3,**dict(zip(keys,values))}
+                accel.print(json.dumps(log),flush=True)
+                if accel.is_main_process:
+                    with (out/'train.jsonl').open('a') as handle:handle.write(json.dumps(log)+'\n')
+            if step%cfg['save_every']==0 or step in cfg['save_steps'] or step==final:save()
+            if step==final:break
+        epoch,consumed=epoch+1,0
+    accel.end_training()
+
+
+if __name__=='__main__':main()

viddub/causal/cf/upstream.py

@@ -0,0 +1,96 @@
+"""Attributed scheduler compatibility from Causal-Forcing.
+
+Adapted from ``utils/scheduler.py`` at commit
+da3ddf1a590f30c9edc568e11c2f77170995c31c. The upstream project is
+Apache-2.0 licensed. Keeping this small compatibility class locally avoids
+requiring users to clone an otherwise unused complete repository.
+"""
+import hashlib
+from pathlib import Path
+
+import torch
+
+COMMIT = "da3ddf1a590f30c9edc568e11c2f77170995c31c"
+REPOSITORY = "https://github.com/thu-ml/Causal-Forcing"
+
+
+def provenance():
+    source = Path(__file__).resolve()
+    return {
+        "repository": REPOSITORY,
+        "commit": COMMIT,
+        "compatibility_file": source.name,
+        "sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
+    }
+
+
+class FlowMatchScheduler:
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000,
+                 shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002,
+                 inverse_timesteps=False, extra_one_step=False,
+                 reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0,
+                      training=False):
+        sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        count = num_inference_steps + 1 if self.extra_one_step else num_inference_steps
+        self.sigmas = torch.linspace(sigma_start, self.sigma_min, count)
+        if self.extra_one_step:
+            self.sigmas = self.sigmas[:-1]
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
+            shifted = y - y.min()
+            self.linear_timesteps_weights = shifted * (num_inference_steps / shifted.sum())
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        shape = (-1,) + (1,) * (sample.ndim - 1)
+        sigma = self.sigmas[timestep_id].reshape(shape)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            next_sigma = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            next_sigma = self.sigmas[timestep_id + 1].reshape(shape)
+        return sample + model_output * (next_sigma - sigma)
+
+    def add_noise(self, original_samples, noise, timestep):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        shape = (-1,) + (1,) * (noise.ndim - 1)
+        sigma = self.sigmas[timestep_id].reshape(shape)
+        return ((1 - sigma) * original_samples + sigma * noise).type_as(noise)
+
+    @staticmethod
+    def training_target(sample, noise, timestep):
+        return noise - sample
+
+    def training_weight(self, timestep):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
+        return self.linear_timesteps_weights[timestep_id]

viddub/causal/infer.py

@@ -0,0 +1,100 @@
+"""Generate one prepared face clip for same/cross audio comparison.
+
+This CLI outputs a cropped clip; full-video detection/paste-back remains in the
+the full-video inference entry point. It does not pretend to be an end-to-end streaming server.
+"""
+import argparse
+import json
+from pathlib import Path
+import time
+
+import numpy as np
+import torch
+
+from .runtime import (atomic_json, autocast, checkpoint_paths, load_components,
+                      local_imports, read_config, require_causal_metadata)
+from .sampling import rollout, schedule
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--config", required=True)
+    p.add_argument("--checkpoint-dir", required=True)
+    p.add_argument("--step", type=int, required=True)
+    p.add_argument("--clip-id", required=True, help="Prepared validation clip id in s02 manifest")
+    p.add_argument("--audio", help="Optional replacement WAV; original clip audio otherwise")
+    p.add_argument("--output", required=True)
+    p.add_argument("--steps", type=int, help="AR defaults to 30; ODE/DMD requires trained step count")
+    p.add_argument("--audio-cfg", type=float, default=1)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--no-cache", action="store_true", help="Reference recomputation implementation")
+    a = p.parse_args()
+    cfg = read_config(a.config)
+    meta = require_causal_metadata(a.checkpoint_dir, a.step, cfg)
+    steps = a.steps if a.steps is not None else (30 if meta["stage"] == "ar" else cfg["student_steps"])
+    if meta["stage"] != "ar":
+        if steps != meta["config"]["student_steps"] or cfg["shift"] != meta["config"]["shift"]:
+            raise ValueError("Few-step inference must use the training schedule")
+        if a.audio_cfg != 1:
+            raise ValueError("Guidance is distilled into ODE/DMD weights; use audio-cfg=1")
+    if Path(a.output).exists():
+        raise FileExistsError(a.output)
+    local_imports()
+    import librosa
+    from viddub.full_context import runtime as full_context_runtime
+    from viddub.full_context.runtime import ZeroLatents, build_inputs, compute_latents, load_clip, load_rows, write_mp4
+    if cfg.get("clips_root"):
+        full_context_runtime.CLIPS = Path(cfg["clips_root"])
+    rows = {r["clip_id"]: r for r in load_rows("val")}
+    if a.clip_id not in rows:
+        raise ValueError("clip-id must be from the validation split")
+    if not torch.cuda.is_available():
+        raise RuntimeError("Full inference requires CUDA")
+    device = torch.device("cuda")
+    pipe, whisper, model = load_components(checkpoint_paths(a.checkpoint_dir, a.step), cfg, device)
+    model.requires_grad_(False).eval()
+    clip = load_clip(rows[a.clip_id], cfg["num_frames"])
+    audio = librosa.load(a.audio, sr=16000, mono=True)[0] if a.audio else clip["audio"]
+    length = round(cfg["num_frames"] / cfg["fps"] * 16000)
+    audio = np.pad(audio[:length], (0, max(0, length - len(audio))))
+    started = time.monotonic()
+    cached = full_context_runtime.CLIPS / rows[a.clip_id]["dir"] / cfg["latents_name"]
+    if cached.is_file():
+        lat = torch.load(cached, map_location="cpu", weights_only=False)
+    else:
+        if pipe.text_encoder is None:
+            raise ValueError("No cached latents: unset VIDDUB_SKIP_TEXT_ENCODER for text encoding")
+        pipe.text_encoder.to(device)
+        lat = compute_latents(pipe, clip)
+        pipe.text_encoder.cpu()
+    x0, vc, ctx = build_inputs(lat, clip["mask"], cfg["num_frames"], ZeroLatents(pipe),
+                                drop_ref=cfg["num_ref"] == 0)
+    dtype = torch.bfloat16 if cfg["precision"] == "bf16" else torch.float32
+    hs = whisper([audio], device, torch.bfloat16)
+    conditions = {"vace_context": vc.to(dtype), "context": ctx.to(dtype),
+                  "whisper_hs": hs.to(dtype), "num_frames": cfg["num_frames"]}
+    torch.manual_seed(a.seed)
+    torch.cuda.synchronize()
+    denoise_start = time.monotonic()
+    with torch.no_grad(), autocast(device, cfg["precision"]):
+        if cfg['sampler']=='cf_renoise':
+            from .cf.sampling import cf_rollout
+            latents=cf_rollout(model,x0.to(dtype),conditions,steps=steps,shift=cfg['shift'],use_cache=not a.no_cache)
+        else:
+            latents = rollout(model, x0.to(dtype), conditions, schedule(steps, cfg["shift"], device),
+                              use_cache=not a.no_cache, audio_cfg=a.audio_cfg)
+    torch.cuda.synchronize()
+    denoise_seconds = time.monotonic() - denoise_start
+    video = pipe.vae.decode(latents[:, :, cfg["num_ref"]:].to(pipe.torch_dtype), device=device)
+    frames = ((video[0].permute(1, 2, 3, 0).float().clamp(-1, 1) + 1) * 127.5).round().byte().cpu().numpy()
+    write_mp4(a.output, frames, audio, cfg["fps"])
+    report = {"clip_id": a.clip_id, "step": a.step, "stage": meta["stage"], "sampling_steps": steps,
+              "cached": not a.no_cache, "sampler":cfg['sampler'], "denoise_seconds": denoise_seconds,
+              "total_seconds": time.monotonic() - started,
+              "frames": len(frames), "note": "Offline clip; excludes model load, no streaming latency claim"}
+    atomic_json(Path(a.output).with_suffix(".json"), report)
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

viddub/causal/make_trajectories.py

@@ -0,0 +1,121 @@
+"""Generate auditable causal teacher rollouts; never initialize from ground-truth video."""
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .runtime import (atomic_json, autocast, checkpoint_paths, fingerprint, load_components,
+                      local_imports, make_dataset, prepare_clip, read_config, require_causal_metadata)
+from .sampling import rollout, schedule
+
+
+def safe_tensor_tree(value):
+    """Convert trajectory payloads to weights_only-safe primitive/tensor trees."""
+    if isinstance(value, np.ndarray):
+        return torch.from_numpy(value)
+    if torch.is_tensor(value):
+        return value.cpu()
+    if isinstance(value, dict):
+        return {k: safe_tensor_tree(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return type(value)(safe_tensor_tree(v) for v in value)
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    raise TypeError(f'Unsafe trajectory payload type: {type(value).__name__}')
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--config", required=True, help="ODE-stage config; schedule must match student training")
+    p.add_argument("--teacher-dir", required=True)
+    p.add_argument("--teacher-step", type=int, required=True)
+    p.add_argument("--output", required=True)
+    p.add_argument("--teacher-steps", type=int, default=32)
+    p.add_argument("--max-clips", type=int, default=256)
+    p.add_argument("--audio-mode", choices=["same", "cross", "both"], default="both")
+    p.add_argument("--shard", type=int, default=0)
+    p.add_argument("--shards", type=int, default=1)
+    p.add_argument("--check-only", action="store_true")
+    a = p.parse_args()
+    cfg = read_config(a.config)
+    paths = checkpoint_paths(a.teacher_dir, a.teacher_step)
+    require_causal_metadata(a.teacher_dir, a.teacher_step, cfg, ("ar",))
+    if a.teacher_steps < cfg["student_steps"] or a.teacher_steps % cfg["student_steps"]:
+        raise ValueError("teacher-steps must be a positive multiple of student_steps")
+    if a.max_clips < 1:
+        raise ValueError("max-clips must be positive")
+    if a.shards < 1 or not 0 <= a.shard < a.shards:
+        raise ValueError("invalid shard assignment")
+    out = Path(a.output)
+    if out.exists() and any(out.iterdir()):
+        raise ValueError("Use a new empty trajectory output directory")
+    dataset, rows = make_dataset(cfg)
+    if a.audio_mode != "same" and len(dataset) < 2:
+        raise ValueError("Cross-audio trajectories need at least two clips")
+    if a.check_only:
+        print(json.dumps({"status": "inputs_validated_no_gpu_execution", "clips": min(a.max_clips, len(dataset)),
+                          "teacher_steps": a.teacher_steps, "student_steps": cfg["student_steps"]}, indent=2))
+        return
+    if not torch.cuda.is_available():
+        raise RuntimeError("Trajectory generation requires CUDA")
+    local_imports()
+    from viddub.full_context import runtime as full_context_runtime
+    from viddub.full_context.runtime import ZeroLatents
+    from .runtime import PixelClipDataset
+    if not isinstance(dataset, PixelClipDataset):
+        dataset = PixelClipDataset(dataset, full_context_runtime.CLIPS)
+    device = torch.device("cuda")
+    pipe, whisper, model = load_components(paths, cfg, device)
+    model.requires_grad_(False).eval()
+    zero = ZeroLatents(pipe)
+    sigmas = schedule(a.teacher_steps, cfg["shift"], device)
+    interval = a.teacher_steps // cfg["student_steps"]
+    manifest = {"schema": 1, "complete": False, "config": cfg,
+                "student_steps": cfg["student_steps"], "teacher_sigmas": sigmas.cpu().tolist(),
+                "boundary_indices": list(range(0, a.teacher_steps + 1, interval)),
+                "teacher": fingerprint(paths), "teacher_cfg": cfg["teacher_cfg"], "records": []}
+    out.mkdir(parents=True, exist_ok=True)
+    atomic_json(out / "manifest.json", manifest)
+    for index in range(a.shard, min(a.max_clips, len(dataset)), a.shards):
+        modes = ("same", "cross") if a.audio_mode == "both" else (a.audio_mode,)
+        for mode in modes:
+            clip = dataset[index]
+            audio_id = clip["clip_id"]
+            if mode == "cross":
+                partner_id = rows[index].get('audio_cross_id')
+                if partner_id is not None:
+                    candidates = [j for j, row in enumerate(rows) if row['clip_id'] == partner_id]
+                    if len(candidates) != 1:
+                        raise ValueError('Frozen cross-audio partner missing or ambiguous')
+                    other = dataset[candidates[0]]
+                else:
+                    other = dataset[(index + 1) % len(dataset)]
+                clip["audio"], audio_id = other["audio"], other["clip_id"]
+            x0, conditions = prepare_clip(clip, pipe, whisper, zero, cfg)
+            conditions['cross'] = mode == 'cross'
+            torch.manual_seed(cfg["seed"] + index * 2 + int(mode == "cross"))
+            with torch.no_grad(), autocast(device, cfg["precision"]):
+                generated, chunks = rollout(model, x0, conditions, sigmas, use_cache=cfg["use_cache"],
+                                            audio_cfg=cfg["teacher_cfg"], trajectory=True)
+            # Persist reference + shape only, never use GT video as trajectory input.
+            shape_ref = torch.zeros_like(x0)
+            shape_ref[:, :, :cfg["num_ref"]] = x0[:, :, :cfg["num_ref"]]
+            name = f"{index:06d}_{mode}.pt"
+            data = {"clip_id": clip["clip_id"], "audio_id": audio_id, "mode": mode,
+                    "shape_reference": shape_ref.cpu(), "gt_latents": x0.cpu(),
+                    "generated": generated.cpu(), "chunks": chunks,
+                    "conditions": safe_tensor_tree(conditions)}
+            torch.save(data, out / (name + ".tmp"))
+            (out / (name + ".tmp")).replace(out / name)
+            manifest["records"].append({"file": name, "clip_id": clip["clip_id"], "audio_id": audio_id,
+                                        "mode": mode, "bytes": (out / name).stat().st_size})
+            atomic_json(out / "manifest.json", manifest)
+            print(json.dumps(manifest["records"][-1]), flush=True)
+    manifest["complete"] = True
+    atomic_json(out / "manifest.json", manifest)
+
+
+if __name__ == "__main__":
+    main()

viddub/causal/model.py

@@ -0,0 +1,304 @@
+"""Causal forward over the retained Wan/VACE and VidDub audio weights.
+
+Parameter names remain compatible with the released checkpoints. The reference
+is a separate prefix block; video blocks attend to themselves and earlier blocks.
+Cache writes are explicit, inference-only transactions on clean history.
+"""
+from dataclasses import asdict, dataclass, field
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+@dataclass(frozen=True)
+class CausalConfig:
+    chunk_size: int = 3  # latent frames; first video chunk includes the VAE first frame
+    num_ref: int = 1
+    fps: float = 24.0
+
+    def __post_init__(self):
+        if self.chunk_size < 1 or self.num_ref not in (0, 1) or self.fps <= 0:
+            raise ValueError("Require chunk_size >= 1, num_ref in {0,1}, fps > 0")
+
+    def blocks(self, frames):
+        if frames <= self.num_ref:
+            raise ValueError("Sequence must contain at least one video latent")
+        out = [(0, self.num_ref)] if self.num_ref else []
+        out += [(i, min(i + self.chunk_size, frames))
+                for i in range(self.num_ref, frames, self.chunk_size)]
+        return out
+
+    def video_blocks(self, frames):
+        return [(a, b) for a, b in self.blocks(frames) if a >= self.num_ref]
+
+
+@dataclass
+class HistoryCache:
+    """Use a NEW instance for every sample/condition/CFG branch/model revision.
+
+    Keys are layer-specific; positions are absolute and never reset at a chunk.
+    No eviction yet: this reference implementation is bounded by training length.
+    """
+    frames: int = 0
+    entries: dict = field(default_factory=dict)
+    geometry: tuple | None = None
+    model_id: int | None = None
+    # None preserves the correctness-reference behavior (all clean history).
+    # Otherwise keep the separate reference prefix plus at most this many
+    # most-recent video latent frames. Absolute RoPE is already embedded in K.
+    max_video_history_frames: int | None = None
+    retained_frames: int = 0
+
+    def __post_init__(self):
+        if self.max_video_history_frames is not None and self.max_video_history_frames < 0:
+            raise ValueError("max_video_history_frames must be nonnegative or None")
+
+
+def _sinusoidal(dim, t):
+    freq = torch.pow(10000, -torch.arange(dim // 2, device=t.device,
+                                       dtype=torch.float64) / (dim // 2))
+    angles = t.flatten().double()[:, None] * freq[None]
+    return torch.cat([angles.cos(), angles.sin()], -1).to(t.dtype)
+
+
+def _rope(x, freqs, heads):
+    b, n, dim = x.shape
+    z = x.reshape(b, n, heads, -1, 2).double().contiguous()
+    z = torch.view_as_complex(z) * freqs
+    return torch.view_as_real(z).flatten(2).to(x.dtype)
+
+
+def _attention(q, k, v, heads):
+    # RMSNorm can promote q/k when trainable master weights are float32.
+    q, k = q.to(v.dtype), k.to(v.dtype)
+    if v.is_cuda and v.dtype in (torch.bfloat16, torch.float16):
+        # Use the SAME fused backend as the SFT teacher. BF16 SDPA vs FA2
+        # differences compound substantially through this trained checkpoint.
+        from viddub.foundation.models.wan_video_dit import flash_attention
+        return flash_attention(q, k, v, heads)
+    def split(z):
+        return z.unflatten(-1, (heads, -1)).transpose(1, 2)
+    y = F.scaled_dot_product_attention(split(q), split(k), split(v), dropout_p=0.0)
+    return y.transpose(1, 2).flatten(2)
+
+
+class CausalWanDub(nn.Module):
+    def __init__(self, dit, vace, adapter, config=CausalConfig(), causal=True):
+        super().__init__()
+        self.dit, self.vace, self.adapter = dit, vace, adapter
+        self.config, self.causal = config, causal
+        if dit.patch_size[0] != 1 or vace.vace_patch_embedding.kernel_size[0] != 1:
+            raise ValueError("Only temporal patch size 1 is supported")
+        if len(adapter.blocks) != len(dit.blocks):
+            raise ValueError("Audio/DiT layer counts differ")
+
+    def _self_attention(self, attn, x, freqs, blocks, tokens_per_frame,
+                        start_frame, cache, pending, key, detach_history_before):
+        q = _rope(attn.norm_q(attn.q(x)), freqs, attn.num_heads)
+        k = _rope(attn.norm_k(attn.k(x)), freqs, attn.num_heads)
+        v = attn.v(x)
+        k = k.to(v.dtype)
+        if cache is None and detach_history_before:
+            stop = detach_history_before * tokens_per_frame
+            k = torch.cat([k[:, :stop].detach(), k[:, stop:]], 1)
+            v = torch.cat([v[:, :stop].detach(), v[:, stop:]], 1)
+        if cache is not None:
+            old = cache.entries.get(key)
+            if cache.frames and old is None:
+                raise ValueError(f"Missing history for {key}")
+            if old is not None:
+                k = torch.cat([old[0], k], 1)
+                v = torch.cat([old[1], v], 1)
+            if pending is not None:
+                pending[key] = (k.detach(), v.detach())
+        if not self.causal:
+            y = _attention(q, k, v, attn.num_heads)
+        elif cache is not None:
+            # A cached call is exactly one complete current block. Its K/V are
+            # retained clean history followed by all tokens of this block, so
+            # every current query may read the entire tensor. Absolute temporal
+            # positions were already applied by RoPE before K entered the cache.
+            y = _attention(q, k, v, attn.num_heads)
+        else:
+            pieces = []
+            for a, b in blocks:
+                a, b = max(a, start_frame), min(b, start_frame + x.shape[1] // tokens_per_frame)
+                if a >= b:
+                    continue
+                qa, qb = (a - start_frame) * tokens_per_frame, (b - start_frame) * tokens_per_frame
+                # All spatial tokens WITHIN this block are mutually visible.
+                pieces.append(_attention(q[:, qa:qb], k[:, :b * tokens_per_frame],
+                                         v[:, :b * tokens_per_frame], attn.num_heads))
+            y = torch.cat(pieces, 1)
+        return attn.o(y)
+
+    def _block(self, block, x, context, t_mod, freqs, blocks, spatial,
+               start, cache, pending, key, detach_history_before):
+        shift, scale, gate, shift_mlp, scale_mlp, gate_mlp = (
+            block.modulation.to(t_mod) + t_mod).unbind(2)
+        h = block.norm1(x) * (1 + scale) + shift
+        x = x + gate * self._self_attention(block.self_attn, h, freqs, blocks,
+                                           spatial, start, cache, pending, key, detach_history_before)
+        ca = block.cross_attn
+        if ca.has_image_input:
+            raise ValueError("Use the VACE reference prefix, not CLIP image conditioning")
+        h = block.norm3(x)
+        x = x + ca.o(_attention(ca.norm_q(ca.q(h)), ca.norm_k(ca.k(context)),
+                                ca.v(context), ca.num_heads))
+        return x + gate_mlp * block.ffn(block.norm2(x) * (1 + scale_mlp) + shift_mlp)
+
+    def forward(self, latents, timesteps, context, vace_context, whisper_hs,
+                num_frames, drop_audio=None, *, start_frame=0, cache=None,
+                commit_cache=False, gradient_checkpointing=False, detach_history_before=0,
+                return_hidden_layer=None, audio_tokens=None):
+        """timesteps: scalar, [B], or [B,F], using t=1000*sigma.
+
+        A cached call receives ONLY a single new block of latents/conditions.
+        whisper_hs remains the full (offline) clip feature tensor, indexed with
+        absolute frame offsets. Reference latents must be clean at sampling time.
+        """
+        b, _, frames, height, width = latents.shape
+        if return_hidden_layer is not None and not 0 <= return_hidden_layer < len(self.dit.blocks):
+            raise ValueError('Invalid hidden layer')
+        total = (num_frames - 1) // 4 + 1 + self.config.num_ref
+        if num_frames < 1 or (num_frames - 1) % 4 or start_frame + frames > total:
+            raise ValueError("num_frames must be 4n+1 and cover this latent slice")
+        if vace_context.shape[0] != b or vace_context.shape[2:] != latents.shape[2:]:
+            raise ValueError("VACE and video latent shapes must match")
+        if height % self.dit.patch_size[1] or width % self.dit.patch_size[2]:
+            raise ValueError("Latent spatial dimensions must be divisible by patch size")
+        blocks = self.config.blocks(total)
+        read_cache = cache
+        if cache is None:
+            if start_frame or commit_cache:
+                raise ValueError("A nonzero start/commit requires a history cache")
+            pending = None
+        else:
+            if not self.causal or start_frame != cache.frames:
+                raise ValueError("Cache must be contiguous and causal")
+            if (start_frame, start_frame + frames) not in blocks:
+                raise ValueError("Cached call must cover exactly one complete block")
+            geometry = (b, height, width, total, self.config, latents.device, latents.dtype)
+            if cache.geometry is not None and cache.geometry != geometry:
+                raise ValueError("Cache geometry changed; create a new cache")
+            if cache.model_id not in (None, id(self)):
+                raise ValueError("Cache belongs to a different model")
+            if commit_cache and (torch.is_grad_enabled() or self.training):
+                raise ValueError("Cache commits require eval() and torch.no_grad()")
+            pending = {} if commit_cache else None
+            # Checkpoint recomputation must see the history from THIS forward,
+            # even after rollout commits subsequent chunks to the live cache.
+            read_cache = HistoryCache(cache.frames, cache.entries, cache.geometry, cache.model_id,
+                                      cache.max_video_history_frames, cache.retained_frames)
+        ts = torch.as_tensor(timesteps, device=latents.device, dtype=latents.dtype)
+        if ts.numel() == 1:
+            ts = ts.reshape(1, 1).expand(b, frames)
+        elif ts.shape == (b,):
+            ts = ts[:, None].expand(b, frames)
+        elif ts.shape != (b, frames):
+            raise ValueError("timesteps must be scalar, [B], or [B,F]")
+        if commit_cache and torch.count_nonzero(ts).item():
+            raise ValueError("Only t=0 clean history may be committed")
+        if audio_tokens is None:
+            tokens = self.adapter.tokens(whisper_hs, num_frames, self.config.num_ref,
+                                         torch.zeros(b, device=latents.device, dtype=torch.bool)
+                                         if drop_audio is None else drop_audio,
+                                         fps=self.config.fps)
+        else:
+            if drop_audio is not None and bool(torch.as_tensor(drop_audio).any()):
+                raise ValueError("Precomputed conditional audio tokens cannot serve a null-audio branch")
+            expected = (num_frames - 1) // 4 + 1 + self.config.num_ref
+            if audio_tokens.shape[:2] != (b, expected):
+                raise ValueError("Precomputed audio-token shape does not match the sequence")
+            tokens = audio_tokens
+        tokens = tokens[:, start_frame:start_frame + frames]
+        x = self.dit.patchify(latents)
+        _, _, f, h, w = x.shape
+        spatial = h * w
+        x = x.flatten(2).transpose(1, 2).contiguous()
+        context = self.dit.text_embedding(context)
+        # Project each distinct timestep before spatial expansion. Projecting
+        # thousands of repeated BF16 rows selects different GEMM kernels than
+        # the original single-timestep path and needlessly amplifies rounding.
+        unique_ts, inverse = torch.unique(ts, return_inverse=True)
+        time_values, mod_values = [], []
+        for value in unique_ts.unbind():
+            emb = self.dit.time_embedding(_sinusoidal(self.dit.freq_dim, value[None]))
+            time_values.append(emb)
+            mod_values.append(self.dit.time_projection(emb))
+        t = torch.cat(time_values)[inverse].reshape(b, f, -1).repeat_interleave(spatial, 1)
+        t_mod = torch.cat(mod_values)[inverse].reshape(b, f, 6, self.dit.dim).repeat_interleave(spatial, 1)
+        ff, fh, fw = self.dit.freqs
+        if start_frame + f > len(ff) or h > len(fh) or w > len(fw):
+            raise ValueError("Position exceeds the model RoPE table")
+        freqs = torch.cat([
+            ff[start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            fh[:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            fw[:w].view(1, 1, w, -1).expand(f, h, w, -1),
+        ], -1).reshape(f * h * w, 1, -1).to(x.device)
+
+        def run(block, value, key):
+            def fn(value):
+                return self._block(block, value, context, t_mod, freqs, blocks,
+                                   spatial, start_frame, read_cache, pending, key, detach_history_before)
+            if gradient_checkpointing and torch.is_grad_enabled():
+                return checkpoint(fn, value, use_reentrant=False)
+            return fn(value)
+
+        # Original VACE concatenates per-example tokens, yielding contiguous
+        # [B,N,C] storage before before_proj. Preserve that GEMM input layout.
+        c = self.vace.vace_patch_embedding(vace_context).flatten(2).transpose(1, 2).contiguous()
+        hints = []
+        for i, block in enumerate(self.vace.vace_blocks):
+            if i == 0:
+                c = block.before_proj(c) + x
+            c = run(block, c, ("vace", i))
+            hints.append(block.after_proj(c))
+        selected_hidden = None
+        for i, block in enumerate(self.dit.blocks):
+            x = run(block, x, ("dit", i))
+            if i in self.vace.vace_layers_mapping:
+                x = x + hints[self.vace.vace_layers_mapping[i]]
+            audio_block = self.adapter.blocks[i]
+            x = checkpoint(audio_block, x, tokens, f, use_reentrant=False) if (
+                gradient_checkpointing and torch.is_grad_enabled()) else audio_block(x, tokens, f)
+            if i == return_hidden_layer:
+                selected_hidden = x.reshape(b, f, h, w, -1)
+        out = self.dit.unpatchify(self.dit.head(x, t), (f, h, w))
+        if pending is not None:
+            # Publish only after EVERY layer and output succeeded.
+            retained_frames = start_frame + frames
+            if cache.max_video_history_frames is not None:
+                keep_video = min(cache.max_video_history_frames,
+                                 max(0, retained_frames - self.config.num_ref))
+                keep_total = self.config.num_ref + keep_video
+                for key, (k, v) in pending.items():
+                    ref_tokens = self.config.num_ref * spatial
+                    video_tokens = keep_video * spatial
+                    if ref_tokens:
+                        k_ref, v_ref = k[:, :ref_tokens], v[:, :ref_tokens]
+                    else:
+                        k_ref, v_ref = k[:, :0], v[:, :0]
+                    if video_tokens:
+                        k_video, v_video = k[:, -video_tokens:], v[:, -video_tokens:]
+                    else:
+                        k_video, v_video = k[:, :0], v[:, :0]
+                    pending[key] = (torch.cat([k_ref, k_video], 1).detach(),
+                                    torch.cat([v_ref, v_video], 1).detach())
+                retained_frames = keep_total
+            cache.entries = pending
+            cache.frames = start_frame + frames
+            cache.retained_frames = retained_frames
+            cache.geometry, cache.model_id = geometry, id(self)
+        return (out, selected_hidden) if return_hidden_layer is not None else out
+
+    def architecture(self):
+        return {"config": asdict(self.config), "causal": self.causal,
+                "dim": self.dit.dim, "blocks": len(self.dit.blocks),
+                "vace_blocks": len(self.vace.vace_blocks),
+                "reference_policy": ("clean_separate_prefix" if self.causal
+                                     else "jointly_noised_full_context"),
+                "history_policy": ("clean_t0" if self.causal else "none"),
+                "audio_policy": "offline_whisper"}

viddub/causal/pretrain_fake.py

@@ -0,0 +1,118 @@
+"""Pretrain and export one reusable fake-score model on frozen student rollouts."""
+import argparse
+import copy
+import datetime
+import json
+import random
+import time
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from .training.stages import FakeObjective
+from .runtime import (atomic_json, autocast, checkpoint_paths, export_weights, fingerprint,
+                      first_item, load_components, make_dataset, prepare_clip, read_config,
+                      require_causal_metadata)
+from .sampling import rollout, schedule
+from .training.engine import check_loss, finite_update, lr_scheduler, optimizer
+
+
+def main():
+    p=argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--config',required=True)
+    p.add_argument('--student-dir',required=True);p.add_argument('--student-step',type=int,required=True)
+    p.add_argument('--teacher-dir',required=True);p.add_argument('--teacher-step',type=int,required=True)
+    p.add_argument('--output',required=True);p.add_argument('--resume')
+    p.add_argument('--stop-after-step',type=int)
+    a=p.parse_args();cfg=read_config(a.config)
+    run_until = cfg['max_steps'] if a.stop_after_step is None else a.stop_after_step
+    if not 0 < run_until <= cfg['max_steps']:
+        raise ValueError('stop-after-step must lie within the configured training budget')
+    if cfg['sampler'] != 'euler':
+        raise ValueError('Shared fake pretraining currently supports only the Euler sampler')
+    if cfg['stage']!='dmd':raise ValueError('Fake pretraining requires a DMD config')
+    student_paths=checkpoint_paths(a.student_dir,a.student_step)
+    require_causal_metadata(a.student_dir,a.student_step,cfg,('ode','dmd'))
+    teacher_paths=checkpoint_paths(a.teacher_dir,a.teacher_step)
+    if a.resume and Path(a.resume).resolve().parent != Path(a.output).resolve():
+        raise ValueError('Resume state must belong to this output directory')
+    if not a.resume and Path(a.output).exists() and any(Path(a.output).iterdir()):
+        raise ValueError('Output is not empty; choose a new output or explicitly --resume')
+    dataset,rows=make_dataset(cfg)
+    from .runtime import local_imports
+    local_imports()
+    import accelerate
+    from viddub.full_context.runtime import ZeroLatents
+    accel=accelerate.Accelerator(mixed_precision='no',gradient_accumulation_steps=1,
+        kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=False,
+            gradient_as_bucket_view=True),accelerate.InitProcessGroupKwargs(timeout=datetime.timedelta(hours=2))])
+    pipe,whisper,student=load_components(student_paths,cfg,accel.device)
+    fake=copy.deepcopy(student)
+    for name,path in teacher_paths.items():
+        getattr(fake,name).load_state_dict(load_file(str(path)),strict=True)
+    student.requires_grad_(False).eval();fake.causal=False;fake.requires_grad_(True).train()
+    objective=FakeObjective(fake,mouth_weight=cfg['mouth_weight'],shift=cfg['shift'],
+        student_steps=cfg['student_steps'],use_cache=cfg['use_cache'],
+        gradient_checkpointing=cfg['gradient_checkpointing'],
+        matching_min=cfg['matching_min'], matching_max=cfg['matching_max'])
+    opt=optimizer(fake,dict(cfg,lr=cfg['fake_lr'],lr_base=cfg['fake_lr_base']))
+    sched=lr_scheduler(opt,cfg)
+    loader=torch.utils.data.DataLoader(dataset,batch_size=1,shuffle=True,collate_fn=first_item,
+        num_workers=cfg['workers'],generator=torch.Generator().manual_seed(cfg['seed']))
+    objective,opt,loader=accel.prepare(objective,opt,loader);accel.register_for_checkpointing(sched)
+    provenance={'student':fingerprint(student_paths),'teacher_init':fingerprint(teacher_paths)}
+    out=Path(a.output);out.mkdir(parents=True,exist_ok=True)
+    if accel.is_main_process:atomic_json(out/'run.json',{'config':cfg,'provenance':provenance,
+        'world_size':accel.num_processes,'kind':'shared_fake_pretrain'})
+    step=epoch=consumed=0
+    random.seed(cfg['seed']+accel.process_index)
+    torch.manual_seed(cfg['seed']+accel.process_index)
+    if a.resume:
+        meta=json.loads((Path(a.resume)/'complete.json').read_text())
+        if meta['provenance']!=provenance or meta['config']!=cfg:raise ValueError('Resume mismatch')
+        if meta['world_size']!=accel.num_processes:raise ValueError('Resume requires same world size')
+        accel.load_state(a.resume)
+        step,epoch,consumed=(meta[k] for k in ('step','epoch','consumed_batches'))
+    zero=ZeroLatents(pipe);start=time.monotonic()
+    row_index={r['clip_id']:i for i,r in enumerate(rows)}
+    while step<run_until:
+        if hasattr(loader,'set_epoch'):loader.set_epoch(epoch)
+        if getattr(loader,'generator',None) is not None:
+            loader.generator.manual_seed(cfg['seed']+epoch)
+        active_loader=accel.skip_first_batches(loader,consumed) if consumed else loader
+        for item in active_loader:
+            consumed+=1
+            if random.random()<cfg['p_cross']:
+                other=dataset[random.randrange(len(dataset))]
+                if other['clip_id']==item['clip_id']:
+                    other=dataset[(row_index[item['clip_id']]+1)%len(dataset)]
+                item=dict(item,audio=other['audio'],cross=True)
+            x0,c=prepare_clip(item,pipe,whisper,zero,cfg)
+            c['cross']=bool(item.get('cross',False))
+            with torch.no_grad(),autocast(accel.device,cfg['precision']):
+                generated=rollout(student,x0,c,schedule(cfg['student_steps'],cfg['shift'],x0.device),
+                                  use_cache=cfg['use_cache'])
+            with autocast(accel.device,cfg['precision']):loss,metrics=objective(x0,c,generated=generated)
+            check_loss(accel,loss);accel.backward(loss);norm=finite_update(accel,objective,opt,cfg);sched.step()
+            step+=1
+            if step%cfg['log_every']==0:accel.print(json.dumps({'step':step,'loss':float(loss.detach()),
+                'grad_norm':norm,'elapsed_seconds':time.monotonic()-start,
+                **{k:float(v) for k,v in metrics.items()}}),flush=True)
+            if step%cfg['save_every']==0 or step==cfg['max_steps'] or step==run_until:
+                state=out/f'state_step{step:06d}'
+                exists=accel.reduce(torch.tensor(int(state.exists()),device=accel.device),'sum')
+                if int(exists):raise FileExistsError(state)
+                accel.save_state(str(state))
+                if accel.is_main_process:
+                    export_weights(accel.unwrap_model(objective).model,out,step,cfg,provenance)
+                    atomic_json(state/'complete.json',{'step':step,'epoch':epoch,
+                        'consumed_batches':consumed,'config':cfg,'provenance':provenance,
+                        'world_size':accel.num_processes})
+                accel.wait_for_everyone()
+            if step==run_until:break
+        epoch,consumed=epoch+1,0
+    accel.end_training()
+
+
+if __name__=='__main__':main()

viddub/causal/quality.py

@@ -0,0 +1,167 @@
+"""Paired GT quality and audio supervision on a causal, trailing decoded window.
+
+Frozen encoders are external to the trainable generator. Decode the complete
+available temporal prefix, so an arbitrary latent slice is never treated as a
+new VAE first frame. Spatial crop is shared by prediction and pixel GT.
+"""
+import torch
+import numpy as np
+import os
+from pathlib import Path
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+def sync_mel_window(mel, start_seconds):
+    from viddub.full_context.sync_loss import mel_window
+    if torch.is_tensor(mel):
+        mel = mel.detach().cpu().numpy()
+    if not isinstance(mel, np.ndarray):
+        raise TypeError(f'Expected mel tensor/array, got {type(mel).__name__}')
+    return mel_window(mel, start_seconds)
+
+
+def trailing_window(latent_frames, size=16):
+    frames = 4 * (latent_frames - 1) + 1
+    return max(0, frames - size), frames
+
+
+def roi_from_mask(mask, latent_height, latent_width, margin=2, max_h=36, max_w=48):
+    ys, xs = torch.where(mask)
+    if not len(ys):
+        raise ValueError('Empty generation mask')
+    t, b, l, r = int(ys.min()), int(ys.max()) + 1, int(xs.min()), int(xs.max()) + 1
+    lt, lb = max(0, t // 8 - margin), min(latent_height, (b + 7) // 8 + margin)
+    ll, lr = max(0, l // 8 - margin), min(latent_width, (r + 7) // 8 + margin)
+    if lb - lt > max_h:
+        lt = (lt + lb - max_h) // 2; lb = lt + max_h
+    if lr - ll > max_w:
+        ll = (ll + lr - max_w) // 2; lr = ll + max_w
+    return lt, lb, ll, lr
+
+
+class GTQuality:
+    def __init__(self, pipe, cfg, device):
+        from .runtime import local_imports
+        local_imports()
+        self.pipe, self.cfg = pipe, cfg
+        self.sync = self.lpips = None
+        if cfg.get('sync_weight', 0) or cfg.get('art_weight', 0):
+            from viddub.full_context import sync_loss
+            from .runtime import PROJECT
+            sync_loss.LS_ROOT = Path(os.environ.get(
+                'VIDDUB_LATENTSYNC_ROOT', PROJECT / 'weights' / 'LatentSync'))
+            self.sync = sync_loss.SyncLoss().to(device).eval().requires_grad_(False)
+        if cfg.get('lpips_weight', 0):
+            from viddub.full_context.runtime import PROJECT
+            import lpips
+            torch.hub.set_dir(str(Path(os.environ.get(
+                'VIDDUB_TORCH_HUB', PROJECT / 'weights' / 'torch_hub'))))
+            self.lpips = lpips.LPIPS(net='vgg', verbose=False).to(device).eval().requires_grad_(False)
+
+    @torch.no_grad()
+    def art_target(self, conditions, end):
+        first,last=trailing_window(end-conditions['num_ref'])
+        if last-first!=16:return None
+        device=conditions['vace_context'].device
+        mask=conditions['pixel_mask'].to(device).bool()
+        ys,xs=torch.where(mask)
+        box=(int(ys.min()),int(ys.max())+1,int(xs.min()),int(xs.max())+1)
+        if conditions.get('cross',False):
+            mel=sync_mel_window(conditions['mel'],first/self.cfg['fps']).to(device)
+            return self.sync.net.audio_encoder(mel).flatten(1).float()
+        gt=conditions['gt_frames'][first:last].to(device).float().permute(3,0,1,2)/127.5-1
+        return self.sync.net.visual_encoder(self.sync.crop(gt,box)).flatten(1).float()
+
+    def decode(self, z):
+        vae = self.pipe.vae
+        # Checkpoint forward and recompute each own a fresh VAE feature cache.
+        vae.model.clear_cache()
+        out = vae.model.decode(z.to(self.pipe.torch_dtype), vae.scale)
+        vae.model.clear_cache()
+        return out
+
+    def decode_window(self, z, grad_begin, first, last):
+        """Exact prefix decode, but frozen GT-history latents require no graph.
+
+        The original conv2 has temporal kernel 1. Processing history under
+        no_grad is equivalent to decoding a detached history with frozen VAE
+        weights, and avoids retaining dozens of irrelevant frame graphs.
+        """
+        vae = self.pipe.vae
+        net = vae.model
+        if net.conv2.kernel_size[0] != 1:
+            raise ValueError('Window decode assumes temporal conv2 kernel 1')
+        net.clear_cache()
+        z = z.to(self.pipe.torch_dtype)
+        mean, inv_std = (s.to(z) for s in vae.scale)
+        values = []
+        grad_enabled = torch.is_grad_enabled()
+        for i in range(z.shape[2]):
+            with torch.set_grad_enabled(grad_enabled and i >= grad_begin):
+                latent = z[:, :, i:i+1] / inv_std.view(1,-1,1,1,1) + mean.view(1,-1,1,1,1)
+                x = net.conv2(latent)
+                decoded, net._feat_map, net._conv_idx = net.decoder(
+                    x, feat_cache=net._feat_map, feat_idx=[0])
+                frame_start = 0 if i == 0 else 4*i-3
+                frame_end = frame_start + decoded.shape[2]
+                lo, hi = max(first, frame_start), min(last, frame_end)
+                if lo < hi:values.append(decoded[:, :, lo-frame_start:hi-frame_start])
+        net.clear_cache()
+        return torch.cat(values, 2)
+
+    def __call__(self, estimate_prefix, conditions, start, end, sigma):
+        zero = estimate_prefix.new_zeros((), dtype=torch.float32)
+        stats = {'pixel': zero, 'lpips': zero, 'sync': zero, 'keep': zero,
+                 'aux_active': zero, 'sync_active': zero}
+        q_active = float(sigma) <= self.cfg['aux_sigma_max']
+        sync_active = float(sigma) <= self.cfg['sync_sigma_max']
+        if not q_active and not sync_active:
+            return zero, stats
+        if conditions.get('gt_frames') is None:
+            raise ValueError('GT auxiliary loss requires original pixel frames')
+        n_ref = conditions['num_ref']
+        z = estimate_prefix[:, :, n_ref:end]
+        if z.shape[2] < 1:return zero, stats
+        mask = conditions['pixel_mask'].to(z.device).bool()
+        lt, lb, ll, lr = roi_from_mask(mask, z.shape[-2], z.shape[-1],
+                                      max_h=self.cfg['aux_max_latent_h'],
+                                      max_w=self.cfg['aux_max_latent_w'])
+        first, last = trailing_window(z.shape[2])
+        z_roi = z[:, :, :, lt:lb, ll:lr]
+        grad_begin = max(0, start - n_ref)
+        decoded = checkpoint(self.decode_window, z_roi, grad_begin, first, last,
+                             use_reentrant=False)[0].float()
+        gt = conditions['gt_frames'][first:last, lt*8:lb*8, ll*8:lr*8].to(z.device).float()
+        gt = gt.permute(3, 0, 1, 2) / 127.5 - 1
+        if decoded.shape != gt.shape:
+            raise ValueError(f'Decoded/GT window mismatch {decoded.shape} != {gt.shape}')
+        local_mask = mask[lt*8:lb*8, ll*8:lr*8]
+        # Q applies only where paired GT exists; current first-wave AR is same.
+        if q_active and not conditions.get('cross', False):
+            if self.cfg.get('pixel_weight', 0):
+                stats['pixel'] = (decoded - gt).abs()[:, :, local_mask].mean()
+            if self.lpips is not None:
+                # Four evenly spaced frames keep memory bounded, including end.
+                ids = torch.linspace(0, last-first-1, min(4, last-first), device=z.device).round().long()
+                ys, xs = torch.where(local_mask)
+                box = (int(ys.min()), int(ys.max())+1, int(xs.min()), int(xs.max())+1)
+                t,b,l,r = box
+                pred = decoded[:, ids, t:b, l:r].transpose(0,1)
+                target = gt[:, ids, t:b, l:r].transpose(0,1)
+                with torch.autocast(z.device.type, enabled=False):
+                    stats['lpips'] = self.lpips(pred.clamp(-1,1), target).mean()
+        if q_active and self.cfg.get('keep_weight', 0) and (~local_mask).any():
+            stats['keep'] = (decoded - gt).abs()[:, :, ~local_mask].mean()
+        dropped = conditions.get('drop_audio')
+        if (sync_active and self.sync is not None and last-first == 16 and
+                not (dropped is not None and bool(dropped.any()))):
+            ys, xs = torch.where(local_mask)
+            box = (int(ys.min()), int(ys.max())+1, int(xs.min()), int(xs.max())+1)
+            mel = sync_mel_window(conditions['mel'], first / self.cfg['fps']).to(z.device)
+            stats['sync'] = self.sync(decoded, box, mel)
+            stats['sync_active'] = zero + 1
+        stats['aux_active'] = zero + float(q_active)
+        loss = sum(self.cfg.get(key+'_weight', 0) * stats[key] for key in ['pixel','lpips','sync','keep'])
+        return loss, stats

viddub/causal/runtime.py

@@ -0,0 +1,343 @@
+"""Local Wan-Dub data/checkpoint integration; imports heavy dependencies lazily."""
+from contextlib import nullcontext
+import hashlib
+import json
+import os
+from pathlib import Path
+
+import torch
+
+from .model import CausalConfig, CausalWanDub
+
+PROJECT = Path(__file__).resolve().parents[2]
+
+
+def local_imports():
+    """Compatibility hook retained for callers; package imports need no path edits."""
+
+
+def atomic_json(path, data):
+    path = Path(path)
+    temp = path.with_name(path.name + ".tmp")
+    temp.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n")
+    temp.replace(path)
+
+
+def read_config(path):
+    # Normalize stored configs as well, so adding DISABLED optional features
+    # does not prevent resuming an older, otherwise identical training state.
+    cfg = dict(path) if isinstance(path, dict) else json.loads(Path(path).read_text())
+    allowed = {"stage", "causal", "chunk_size", "num_ref", "fps", "num_frames", "max_steps", "warmup",
+               "lr", "lr_base", "weight_decay", "grad_accum", "grad_clip", "precision", "workers",
+               "seed", "max_train_clips", "latents_name", "save_every", "log_every", "mouth_weight",
+               "p_drop_audio", "p_cross", "shift", "student_steps", "teacher_cfg", "use_cache",
+               "gradient_checkpointing", "fake_updates", "matching_min", "matching_max",
+               "rows_file", "clips_root", "pixel_weight", "lpips_weight", "sync_weight",
+               "keep_weight", "aux_sigma_max", "sync_sigma_max", "aux_max_latent_h", "aux_max_latent_w",
+               "lr_schedule", "lr_min_ratio", "save_steps", "save_state_every", "ema_decay",
+               "dmd_rollout_mode", "art_weight", "art_layer", "art_warmup", "art_lr"}
+    allowed.update({"ode_gt_flow_weight", "rollout_quality_weight", "rollout_quality_cross",
+                    "dmd_retention_weight"})
+    allowed.update({"fake_lr", "fake_lr_base"})
+    allowed.update({"dmd_score_offload", "dmd_split_quality_backward"})
+    allowed.update({'cd_grid_steps','cd_ema_decay','cd_student_cache','cd_frozen_cache',
+                    'cd_aux_scale','cd_all_chunks','beta1','beta2','sampler'})
+    allowed.update({"ar_history_mode", "ar_history_prob_start", "ar_history_prob_end",
+                    "ar_history_ramp_steps", "ar_history_steps", "ar_history_sigma_max",
+                    "ar_blocks_per_sample", "ar_history_mouth_drop_prob",
+                    "ar_history_mouth_drop_cross_only"})
+    unknown = set(cfg) - allowed
+    if unknown:
+        raise ValueError(f"Unknown config fields: {sorted(unknown)}")
+    defaults = dict(stage="ar", causal=True, chunk_size=3, num_ref=1, fps=24, num_frames=81,
+                    max_steps=2000, warmup=100, lr=1e-5, lr_base=2e-6, weight_decay=0.01,
+                    grad_accum=1, grad_clip=1.0, precision="bf16", workers=2, seed=1234,
+                    max_train_clips=None, latents_name="latents.pt", save_every=250, log_every=10,
+                    mouth_weight=3, p_drop_audio=0.1, p_cross=0.5, shift=5, student_steps=4, teacher_cfg=1,
+                    use_cache=True, gradient_checkpointing=True, fake_updates=5,
+                    matching_min=0.02, matching_max=0.98, rows_file=None, clips_root=None,
+                    pixel_weight=0.0, lpips_weight=0.0, sync_weight=0.0, keep_weight=0.0,
+                    aux_sigma_max=0.5, sync_sigma_max=None,
+                    aux_max_latent_h=36, aux_max_latent_w=48,
+                    lr_schedule='constant', lr_min_ratio=0.1, save_steps=[], save_state_every=0,
+                    ema_decay=0.0, dmd_rollout_mode='legacy_early_exit',
+                    art_weight=0.0, art_layer=8, art_warmup=50, art_lr=1e-4,
+                    ode_gt_flow_weight=0.0, rollout_quality_weight=1.0,
+                    rollout_quality_cross=False, dmd_retention_weight=0.0,
+                    fake_lr=None, fake_lr_base=None, dmd_score_offload=False,
+                    dmd_split_quality_backward=False)
+    defaults.update(ar_history_mode='gt', ar_history_prob_start=0.0,
+                    ar_history_prob_end=0.0, ar_history_ramp_steps=0,
+                    ar_history_steps=4, ar_history_sigma_max=0.5,
+                    ar_blocks_per_sample=1, ar_history_mouth_drop_prob=0.0,
+                    ar_history_mouth_drop_cross_only=False)
+    defaults.update(cd_grid_steps=48, cd_ema_decay=.99, cd_student_cache=False,
+                    cd_frozen_cache=False, cd_aux_scale=1., cd_all_chunks=True,
+                    beta1=.9, beta2=.999, sampler='euler')
+    defaults.update(cfg)
+    cfg = defaults
+    if cfg['sync_sigma_max'] is None:
+        cfg['sync_sigma_max'] = cfg['aux_sigma_max']
+    CausalConfig(cfg["chunk_size"], cfg["num_ref"], cfg["fps"])
+    if cfg["stage"] not in ("ar", "ode", "dmd", "cd") or cfg["precision"] not in ("bf16", "fp32"):
+        raise ValueError("Unknown stage or precision")
+    if cfg["num_frames"] < 1 or (cfg["num_frames"] - 1) % 4:
+        raise ValueError("num_frames must be 4n+1")
+    if cfg["fps"] != 24:
+        raise ValueError("Existing cached training dataset is 24fps")
+    for name in ("max_steps", "lr", "lr_base", "grad_accum", "grad_clip", "save_every",
+                 "log_every", "shift", "student_steps", "fake_updates"):
+        if cfg[name] <= 0:
+            raise ValueError(f"{name} must be positive")
+    if cfg["workers"] < 0 or cfg["warmup"] < 0 or not 0 <= cfg["p_drop_audio"] <= 1 or not 0 <= cfg["p_cross"] <= 1:
+        raise ValueError("Invalid workers/warmup/audio dropout")
+    if cfg["mouth_weight"] < 1 or cfg["teacher_cfg"] < 0:
+        raise ValueError("Invalid mouth weight or guidance")
+    if not 0 < cfg["matching_min"] < cfg["matching_max"] < 1:
+        raise ValueError("Invalid DMD matching range")
+    if cfg["stage"] == "dmd" and cfg["grad_accum"] != 1:
+        raise ValueError("DMD alternating optimizers currently require grad_accum=1")
+    if any(cfg[k] < 0 for k in ['pixel_weight','lpips_weight','sync_weight','keep_weight',
+                                'ode_gt_flow_weight','rollout_quality_weight','dmd_retention_weight']):
+        raise ValueError('Auxiliary weights must be nonnegative')
+    if (cfg['lr_schedule'] not in ('constant','cosine') or
+            not 0 < cfg['aux_sigma_max'] <= 1 or not 0 < cfg['sync_sigma_max'] <= 1):
+        raise ValueError('Invalid scheduler/aux sigma')
+    if not 0 <= cfg['ema_decay'] < 1:raise ValueError('Invalid EMA decay')
+    if cfg['dmd_rollout_mode'] not in ('legacy_early_exit','final_step'):
+        raise ValueError('Invalid DMD rollout mode')
+    if not isinstance(cfg['dmd_split_quality_backward'], bool):
+        raise ValueError('dmd_split_quality_backward must be boolean')
+    if cfg['dmd_split_quality_backward'] and cfg['stage'] != 'dmd':
+        raise ValueError('dmd_split_quality_backward is DMD-only')
+    if (cfg['stage'] == 'dmd' and any(cfg[k] for k in
+            ['pixel_weight','lpips_weight','sync_weight','keep_weight','dmd_retention_weight'])
+            and cfg['dmd_rollout_mode'] != 'final_step'):
+        raise ValueError('DMD quality/retention losses require final_step rollout')
+    if cfg['art_weight'] < 0 or cfg['art_layer'] < 0 or cfg['art_warmup'] < 0 or cfg['art_lr'] <= 0:
+        raise ValueError('Invalid Art configuration')
+    cfg['fake_lr'] = cfg['lr'] if cfg['fake_lr'] is None else cfg['fake_lr']
+    cfg['fake_lr_base'] = cfg['lr_base'] if cfg['fake_lr_base'] is None else cfg['fake_lr_base']
+    if cfg['fake_lr'] <= 0 or cfg['fake_lr_base'] <= 0:
+        raise ValueError('Fake-score learning rates must be positive')
+    if cfg['ar_history_mode'] not in ('gt','resample','rollout'):
+        raise ValueError('Invalid AR history mode')
+    if not 0 <= cfg['ar_history_prob_start'] <= 1 or not 0 <= cfg['ar_history_prob_end'] <= 1:
+        raise ValueError('AR history probabilities must lie in [0,1]')
+    if cfg['ar_history_ramp_steps'] < 0 or cfg['ar_history_steps'] < 1:
+        raise ValueError('Invalid AR history schedule')
+    if not 0 < cfg['ar_history_sigma_max'] <= 1:
+        raise ValueError('Invalid AR history sigma maximum')
+    if not isinstance(cfg['ar_blocks_per_sample'], int) or cfg['ar_blocks_per_sample'] < 0:
+        raise ValueError('ar_blocks_per_sample must be an integer >= 0 (0 means all blocks)')
+    if cfg['stage'] != 'ar' and cfg['ar_blocks_per_sample'] != 1:
+        raise ValueError('ar_blocks_per_sample is AR-only')
+    if not 0 <= cfg['ar_history_mouth_drop_prob'] <= 1:
+        raise ValueError('ar_history_mouth_drop_prob must lie in [0,1]')
+    if not isinstance(cfg['ar_history_mouth_drop_cross_only'], bool):
+        raise ValueError('ar_history_mouth_drop_cross_only must be boolean')
+    if cfg['stage'] != 'ar' and (cfg['ar_history_mouth_drop_prob'] or
+                                cfg['ar_history_mouth_drop_cross_only']):
+        raise ValueError('AR history mouth dropout is AR-only')
+    if cfg['ema_decay'] != 0:
+        raise ValueError('EMA export is not implemented in this trainer; use ema_decay=0')
+    if cfg['cd_grid_steps'] < 2 or not 0 <= cfg['cd_ema_decay'] < 1 or cfg['cd_aux_scale'] < 0:
+        raise ValueError('Invalid consistency configuration')
+    if not 0 <= cfg['beta1'] < 1 or not 0 <= cfg['beta2'] < 1:
+        raise ValueError('Invalid optimizer betas')
+    if cfg['sampler'] not in ('euler','cf_renoise'):
+        raise ValueError('Invalid sampler')
+    return cfg
+
+
+def checkpoint_paths(directory, step):
+    if not directory or step is None:
+        raise ValueError("Provide a checkpoint directory AND step")
+    paths = {name: Path(directory).resolve() / f"{name}_step{step:06d}.safetensors"
+             for name in ("dit", "vace", "adapter")}
+    for p in paths.values():
+        if not p.is_file():
+            raise FileNotFoundError(p)
+    return paths
+
+
+def fingerprint(paths):
+    result = {}
+    for name, path in paths.items():
+        h = hashlib.sha256()
+        with path.open("rb") as f:
+            for chunk in iter(lambda: f.read(8 * 1024**2), b""):
+                h.update(chunk)
+        result[name] = {"path": str(path), "bytes": path.stat().st_size, "sha256": h.hexdigest()}
+    return result
+
+
+def load_components(paths, cfg, device):
+    local_imports()
+    from safetensors.torch import load_file
+    from viddub.full_context import runtime as full_context_runtime
+    # Explicit experiment-local cache, never changes the old training module
+    # on disk. Other processes retain their existing paths.
+    if os.environ.get("VIDDUB_WEIGHT_ROOT"):
+        full_context_runtime.WEIGHTS = Path(os.environ["VIDDUB_WEIGHT_ROOT"])
+        full_context_runtime.VACE_DIR = full_context_runtime.WEIGHTS / "Wan2.1-VACE-1.3B"
+    from viddub.full_context.audio import DubAudioAdapter
+    # A skipped T5 is safe only with the exact precomputed empty-prompt tensor.
+    # If no cache was supplied, load/encode/unload T5 once instead of producing
+    # a tokenizer=None failure on the first minibatch.
+    empty_path = os.environ.get("VIDDUB_EMPTY_CONTEXT")
+    empty_context = None
+    if empty_path:
+        empty_context = torch.load(empty_path, map_location="cpu", weights_only=True)
+        if (not torch.is_tensor(empty_context) or empty_context.ndim != 2 or
+                empty_context.shape[1] != 4096 or not 1 <= empty_context.shape[0] <= 512 or
+                not torch.isfinite(empty_context).all()):
+            raise ValueError("VIDDUB_EMPTY_CONTEXT must contain finite [L,4096] empty T5 embeddings")
+    if os.environ.get("VIDDUB_SKIP_TEXT_ENCODER") == "1" and empty_context is not None:
+        # Cached-data experiments never invoke T5; avoid loading its 11 GB
+        # checkpoint for every short diagnostic/training run.
+        from viddub.foundation.pipelines.wan_video import WanVideoPipeline, ModelConfig
+        from viddub.full_context.audio import WhisperFeatures
+        pipe = WanVideoPipeline.from_pretrained(
+            torch_dtype=torch.bfloat16, device=device,
+            model_configs=[ModelConfig(path=str(full_context_runtime.VACE_DIR / "diffusion_pytorch_model.safetensors")),
+                           ModelConfig(path=str(full_context_runtime.VACE_DIR / "Wan2.1_VAE.pth"))],
+            tokenizer_config=None)
+        for m in (pipe.dit, pipe.vace, pipe.vae):
+            m.requires_grad_(False).eval()
+        whisper = WhisperFeatures(str(full_context_runtime.WEIGHTS / "whisper-small")).to(device, torch.bfloat16)
+    else:
+        pipe, whisper = full_context_runtime.load_models(device)
+        # Evaluation needs only the empty prompt. Encode it once while T5 is
+        # resident on the correct device, then keep the frozen encoder off GPU.
+        pipe._wan_dub_empty_context = full_context_runtime.encode_text(pipe, "").detach().cpu()
+    if empty_context is not None:
+        pipe._wan_dub_empty_context = empty_context
+    adapter = DubAudioAdapter().to(device)
+    for name, module in (("dit", pipe.dit), ("vace", pipe.vace), ("adapter", adapter)):
+        # Preserve FP32 fine-tuning updates when loading a causal checkpoint;
+        # loading into a BF16 base module first silently rounds those updates.
+        module.float()
+        module.load_state_dict(load_file(str(paths[name])), strict=True)
+    model = CausalWanDub(
+        pipe.dit,
+        pipe.vace,
+        adapter,
+        CausalConfig(cfg["chunk_size"], cfg["num_ref"], cfg["fps"]),
+        causal=cfg.get("causal", True),
+    )
+    # FP32 master weights in both modes; BF16 only applies to compute.
+    model.float()
+    # Training/inference here consumes cached text. Keep T5 off the GPU unless
+    # an explicit encoding operation calls the pipeline's model manager.
+    if pipe.text_encoder is not None:
+        pipe.text_encoder.cpu()
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    return pipe, whisper, model
+
+
+def autocast(device, precision):
+    if precision == "bf16":
+        # Rollout switches no_grad/grad within one autocast scope. Reusing a
+        # weight cast cached during no_grad can silently drop weight gradients
+        # and changes checkpoint-recomputation saved tensor metadata.
+        return torch.autocast(device.type, dtype=torch.bfloat16, cache_enabled=False)
+    return nullcontext()
+
+
+def export_weights(model, directory, step, cfg, provenance):
+    from safetensors.torch import save_file
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    for name in ("dit", "vace", "adapter"):
+        dst = directory / f"{name}_step{step:06d}.safetensors"
+        temp = dst.with_suffix(".safetensors.tmp")
+        # Keep fp32 masters; existing inference loaders may cast to bf16.
+        save_file({k: v.detach().float().cpu().contiguous()
+                   for k, v in getattr(model, name).state_dict().items()}, str(temp))
+        temp.replace(dst)
+    atomic_json(directory / f"causal_step{step:06d}.json",
+                {"schema": 1, "step": step, "stage": cfg["stage"], "config": cfg,
+                 "architecture": model.architecture(), "provenance": provenance,
+                 "quality_validated": False})
+
+
+def require_causal_metadata(directory, step, cfg, stages=None):
+    path = Path(directory) / f"causal_step{step:06d}.json"
+    data = json.loads(path.read_text())
+    if data.get("architecture", {}).get(
+        "causal", data["config"].get("causal", True)
+    ) != cfg["causal"]:
+        raise ValueError("Checkpoint/config causal architecture mismatch")
+    for key in ("chunk_size", "num_ref", "fps", "num_frames", "shift"):
+        if data["config"][key] != cfg[key]:
+            raise ValueError(f"Checkpoint/config mismatch for {key}")
+    if stages and data["stage"] not in stages:
+        raise ValueError(f"Expected stages {stages}, got {data['stage']}")
+    if data["stage"] in ("ode", "dmd", "cd") and data["config"]["student_steps"] != cfg["student_steps"]:
+        raise ValueError("Checkpoint/config mismatch for student_steps")
+    if data['stage'] in ('ode','dmd','cd') and data['config'].get('sampler','euler') != cfg.get('sampler','euler'):
+        raise ValueError('Checkpoint/config sampler mismatch')
+    return data
+
+
+def make_dataset(cfg):
+    local_imports()
+    from viddub.full_context import runtime as full_context_runtime
+    from viddub.full_context.runtime import CachedClipDataset, load_rows, cached_clip_ids
+    if cfg.get("clips_root"):
+        full_context_runtime.CLIPS = Path(cfg["clips_root"])
+    rows = json.loads(Path(cfg["rows_file"]).read_text()) if cfg.get("rows_file") else load_rows("train")
+    sources = sorted({r.get('source', 'h3') for r in rows})
+    ids = cached_clip_ids(sources)
+    if cfg.get('rows_file'):
+        missing = [r['clip_id'] for r in rows if r['clip_id'] not in ids]
+        if missing:
+            raise ValueError(f'Explicit manifest contains {len(missing)} uncached clips: {missing[:5]}')
+    rows = [r for r in rows if r["clip_id"] in ids]
+    if cfg["max_train_clips"] is not None:
+        rows = rows[:cfg["max_train_clips"]]
+    if not rows:
+        raise ValueError("No cached training clips")
+    dataset = CachedClipDataset(rows, cfg["num_frames"], cfg["latents_name"])
+    if any(cfg.get(k, 0) for k in ['pixel_weight','lpips_weight','sync_weight','keep_weight',
+                                   'art_weight']):
+        dataset = PixelClipDataset(dataset, full_context_runtime.CLIPS)
+    return dataset, rows
+
+
+class PixelClipDataset(torch.utils.data.Dataset):
+    def __init__(self, dataset, root):
+        self.dataset, self.root = dataset, Path(root)
+    def __len__(self):return len(self.dataset)
+    def __getitem__(self, index):
+        import decord
+        value = self.dataset[index]
+        path = self.root / self.dataset.rows[index]['dir'] / 'video.mp4'
+        reader = decord.VideoReader(str(path), num_threads=1)
+        value['gt_frames'] = torch.from_numpy(reader.get_batch(list(range(self.dataset.num_frames))).asnumpy())
+        return value
+
+
+def first_item(batch):
+    return batch[0]
+
+
+def prepare_clip(clip, pipe, whisper, zero_latents, cfg):
+    local_imports()
+    from viddub.full_context.runtime import build_inputs
+    x0, vc, ctx = build_inputs(clip["latents"], clip["mask"], cfg["num_frames"], zero_latents,
+                               drop_ref=cfg["num_ref"] == 0)
+    if x0.shape[2] != (cfg["num_frames"] - 1) // 4 + 1 + cfg["num_ref"]:
+        raise ValueError("Cached clip is shorter than configured training length")
+    hs = whisper([clip["audio"]], pipe.device, torch.bfloat16)
+    # Whisper remains frozen BF16; transformer input follows compute precision.
+    dtype = torch.bfloat16 if cfg["precision"] == "bf16" else torch.float32
+    conditions = {"context": ctx.to(dtype), "vace_context": vc.to(dtype),
+                  "whisper_hs": hs.to(dtype), "num_frames": cfg["num_frames"], 'num_ref':cfg['num_ref']}
+    if 'gt_frames' in clip:
+        from viddub.full_context.sync_loss import melspectrogram
+        conditions.update(gt_frames=clip['gt_frames'], pixel_mask=torch.from_numpy(clip['mask']),
+                          mel=melspectrogram(clip['audio']))
+    return x0.to(dtype), conditions

viddub/causal/sampling.py

@@ -0,0 +1,298 @@
+"""Flow schedules and clean-history autoregressive rollout.
+
+This is a correctness reference: no cache eviction and no streaming VAE decoder.
+All denoising passes READ history. One extra t=0 pass commits each final chunk.
+"""
+from contextlib import contextmanager
+
+import torch
+
+from .model import HistoryCache
+
+
+def schedule(steps, shift=5.0, device="cpu"):
+    if steps < 1 or shift <= 0:
+        raise ValueError("steps and shift must be positive")
+    s = torch.linspace(1, 0, steps + 1, device=device, dtype=torch.float32)
+    return shift * s / (1 + (shift - 1) * s)
+
+
+def validate_schedule(sigmas):
+    if sigmas.ndim != 1 or len(sigmas) < 2 or not torch.isfinite(sigmas).all():
+        raise ValueError("Expected finite 1D schedule")
+    if abs(float(sigmas[0]) - 1) > 1e-6 or float(sigmas[-1]) != 0:
+        raise ValueError("Schedule must go from pure noise (1) to clean (0)")
+    if not torch.all(sigmas[:-1] > sigmas[1:]):
+        raise ValueError("Schedule must be strictly decreasing")
+
+
+@contextmanager
+def evaluating(model):
+    old = model.training
+    model.eval()
+    try:
+        yield
+    finally:
+        model.train(old)
+
+
+def predict_chunk(model, chunk, sigma, prefix, conditions, start, cache=None,
+                  drop_audio=None, gradient_checkpointing=False, detach_history=False,
+                  return_hidden_layer=None):
+    """Cache-free and cached implementations have exactly the same semantics."""
+    if cache is not None:
+        vc_start = int(conditions.get("vace_start", 0))
+        local_start = start-vc_start
+        if local_start < 0 or local_start+chunk.shape[2] > conditions["vace_context"].shape[2]:
+            raise ValueError("Current VACE condition does not cover the requested absolute block")
+        x = chunk
+        vc = conditions["vace_context"][:, :, local_start:local_start + chunk.shape[2]]
+        offset = start
+        ts = chunk.new_full((chunk.shape[0], chunk.shape[2]), float(sigma) * 1000)
+    else:
+        x = torch.cat([prefix, chunk], 2)
+        vc, offset = conditions["vace_context"][:, :, :x.shape[2]], 0
+        ts = x.new_zeros((x.shape[0], x.shape[2]))
+        ts[:, start:] = float(sigma) * 1000
+    out = model(x, ts, conditions["context"], vc, conditions["whisper_hs"],
+                conditions["num_frames"], drop_audio, start_frame=offset, cache=cache,
+                gradient_checkpointing=gradient_checkpointing,
+                detach_history_before=start if detach_history and cache is None else 0,
+                return_hidden_layer=return_hidden_layer,
+                audio_tokens=conditions.get("audio_tokens"))
+    if return_hidden_layer is not None:
+        prediction, hidden = out
+        return (prediction if cache is not None else prediction[:, :, start:]), hidden
+    return out if cache is not None else out[:, :, start:]
+
+
+def commit(model, chunk, conditions, start, cache, drop_audio=None):
+    if cache is None:
+        return
+    vc_start = int(conditions.get("vace_start", 0))
+    local_start = start-vc_start
+    if local_start < 0 or local_start+chunk.shape[2] > conditions["vace_context"].shape[2]:
+        raise ValueError("Current VACE condition does not cover the committed absolute block")
+    with torch.no_grad():
+        model(chunk.detach(), 0.0, conditions["context"],
+              conditions["vace_context"][:, :, local_start:local_start + chunk.shape[2]],
+              conditions["whisper_hs"], conditions["num_frames"], drop_audio,
+              start_frame=start, cache=cache, commit_cache=True,
+              audio_tokens=conditions.get("audio_tokens"))
+
+
+def history_chunk(chunk, conditions, start, mode):
+    """Inference-only intervention on what is committed as clean visual history.
+
+    The returned tensor affects future chunks only; the visible output for the
+    current chunk remains unchanged. VACE channels 0:C contain the inactive
+    source latent and channels 32: encode the editable-region mask.
+    """
+    if mode == "full":
+        return chunk.detach()
+    vc_start = int(conditions.get("vace_start", 0))
+    local_start = start-vc_start
+    if local_start < 0 or local_start+chunk.shape[2] > conditions["vace_context"].shape[2]:
+        raise ValueError("Current VACE condition does not cover the history block")
+    inactive = conditions["vace_context"][:, :chunk.shape[1], local_start:local_start + chunk.shape[2]].to(chunk)
+    if mode == "inactive":
+        return inactive.detach()
+    if mode == "mouth_suppressed":
+        mask = conditions["vace_context"][:, 32:, local_start:local_start + chunk.shape[2]].amax(
+            1, keepdim=True).float().clamp(0, 1).to(chunk)
+        return (chunk.detach() * (1 - mask) + inactive * mask).detach()
+    raise ValueError(f"Unknown history mode: {mode}")
+
+
+def rollout(model, shape_reference, conditions, sigmas, *, noise=None,
+            use_cache=True, audio_cfg=1.0, grad_chunk=None, grad_step=None,
+            gradient_checkpointing=False, trajectory=False, history_mode="full",
+            on_chunk=None, max_video_history_frames=None):
+    """Generate all video chunks, retaining gradients for at most one chunk/step.
+
+    Normal inference: grad_chunk=None. DMD: randomly choose a video chunk and
+    an exit step; use its x0 estimate and detach every other rollout operation.
+    Later chunks condition on that estimate, but gradients do not pass through
+    cache/history. This deliberately implements truncated rollout gradients.
+    shape_reference video values are NEVER read; only the clean reference prefix
+    (if configured) is consumed. noise has the full [B,C,F,H,W] shape.
+    """
+    validate_schedule(sigmas)
+    if audio_cfg < 0:
+        raise ValueError("audio_cfg must be nonnegative")
+    if not model.causal:
+        if use_cache:
+            raise ValueError(
+                "Bidirectional full-context rollout cannot use a history cache")
+        if on_chunk is not None or max_video_history_frames is not None:
+            raise ValueError(
+                "Streaming callbacks/history bounds are causal-only")
+        noise = torch.randn_like(shape_reference) if noise is None else noise
+        if noise.shape != shape_reference.shape:
+            raise ValueError("noise shape mismatch")
+        ranges = model.config.video_blocks(shape_reference.shape[2])
+        if grad_chunk is not None and not 0 <= grad_chunk < len(ranges):
+            raise ValueError("grad_chunk is outside video chunks")
+        if grad_step is not None and not 0 <= grad_step < len(sigmas) - 1:
+            raise ValueError("grad_step outside schedule")
+        null = torch.ones(
+            shape_reference.shape[0],
+            device=shape_reference.device,
+            dtype=torch.bool,
+        )
+        current = noise
+        states = [current.detach().cpu()] if trajectory else None
+        stop = grad_step if grad_step is not None else len(sigmas) - 2
+        with evaluating(model):
+            for step in range(stop + 1):
+                keep_grad = grad_chunk is not None and step == stop
+                with torch.set_grad_enabled(keep_grad):
+                    ts = current.new_full(
+                        (current.shape[0], current.shape[2]),
+                        float(sigmas[step]) * 1000,
+                    )
+                    pred = model(
+                        current,
+                        ts,
+                        conditions["context"],
+                        conditions["vace_context"][:, :, :current.shape[2]],
+                        conditions["whisper_hs"],
+                        conditions["num_frames"],
+                        gradient_checkpointing=gradient_checkpointing,
+                        audio_tokens=conditions.get("audio_tokens"),
+                    )
+                    if audio_cfg != 1:
+                        pred_null = model(
+                            current,
+                            ts,
+                            conditions["context"],
+                            conditions["vace_context"][:, :, :current.shape[2]],
+                            conditions["whisper_hs"],
+                            conditions["num_frames"],
+                            null,
+                            gradient_checkpointing=gradient_checkpointing,
+                            audio_tokens=conditions.get("audio_tokens"),
+                        )
+                        pred = pred_null + audio_cfg * (pred - pred_null)
+                    if keep_grad and grad_step is not None:
+                        current = current.float() - sigmas[step] * pred.float()
+                    else:
+                        current = (
+                            current.float()
+                            + (sigmas[step + 1] - sigmas[step]) * pred.float()
+                        )
+                    current = current.to(shape_reference.dtype)
+                if trajectory:
+                    states.append(current.detach().cpu())
+        if trajectory:
+            record = {
+                "start": 0,
+                "end": shape_reference.shape[2],
+                "states": torch.stack(states),
+            }
+            return current, [record]
+        return current
+    if history_mode not in ("full", "mouth_suppressed", "inactive"):
+        raise ValueError("Unsupported history mode")
+    n_ref = model.config.num_ref
+    ranges = model.config.video_blocks(shape_reference.shape[2])
+    if grad_chunk is not None and not 0 <= grad_chunk < len(ranges):
+        raise ValueError("grad_chunk is outside video chunks")
+    if grad_step is not None and not 0 <= grad_step < len(sigmas) - 1:
+        raise ValueError("grad_step outside schedule")
+    noise = torch.randn_like(shape_reference) if noise is None else noise
+    if noise.shape != shape_reference.shape:
+        raise ValueError("noise shape mismatch")
+    cache = HistoryCache(max_video_history_frames=max_video_history_frames) if use_cache else None
+    null_cache = HistoryCache(max_video_history_frames=max_video_history_frames) if use_cache and audio_cfg != 1 else None
+    null = torch.ones(shape_reference.shape[0], device=shape_reference.device, dtype=torch.bool)
+    prefix = shape_reference[:, :, :n_ref].detach()
+    chunks, records = [prefix], []
+    with evaluating(model):
+        if n_ref:
+            commit(model, prefix, conditions, 0, cache)
+            if null_cache is not None:
+                commit(model, prefix, conditions, 0, null_cache, null)
+        for index, (start, end) in enumerate(ranges):
+            current = noise[:, :, start:end]
+            states = [current.detach().cpu()] if trajectory else None
+            stop = grad_step if index == grad_chunk and grad_step is not None else len(sigmas) - 2
+            for step in range(stop + 1):
+                keep_grad = index == grad_chunk and step == stop
+                with torch.set_grad_enabled(keep_grad):
+                    pred = predict_chunk(model, current, sigmas[step], prefix.detach(), conditions,
+                                         start, cache, gradient_checkpointing=gradient_checkpointing,
+                                         detach_history=True)
+                    if audio_cfg != 1:
+                        pred_null = predict_chunk(model, current, sigmas[step], prefix.detach(), conditions,
+                                                  start, null_cache, null,
+                                                  gradient_checkpointing=gradient_checkpointing,
+                                                  detach_history=True)
+                        pred = pred_null + audio_cfg * (pred - pred_null)
+                    if keep_grad and grad_step is not None:
+                        current = current.float() - sigmas[step] * pred.float()
+                    else:
+                        current = current.float() + (sigmas[step + 1] - sigmas[step]) * pred.float()
+                    current = current.to(shape_reference.dtype)
+                if trajectory:
+                    states.append(current.detach().cpu())
+            chunks.append(current)
+            if trajectory:
+                records.append({"start": start, "end": end, "states": torch.stack(states)})
+            committed = history_chunk(current, conditions, start, history_mode)
+            commit(model, committed, conditions, start, cache)
+            if null_cache is not None:
+                commit(model, committed, conditions, start, null_cache, null)
+            prefix = torch.cat([prefix.detach(), committed], 2)
+            if on_chunk is not None:
+                on_chunk(index=index, start=start, end=end, chunk=current.detach())
+    out = torch.cat(chunks, 2)
+    return (out, records) if trajectory else out
+
+
+def streaming_rollout_blocks(model, reference, reference_vace, block_inputs,
+                             conditions, sigmas, *, max_video_history_frames,
+                             history_mode="full"):
+    """Yield causal blocks without materializing full-length video/VACE tensors.
+
+    ``block_inputs`` is an iterable of dictionaries containing absolute ``start``,
+    initial ``noise`` [B,C,F,H,W], and local ``vace_context`` for exactly that
+    block. ``conditions`` contains global text/audio tensors, ``num_frames`` and
+    optional precomputed ``audio_tokens``. The reference and every generated block
+    are committed once to one bounded clean-history cache.
+    """
+    validate_schedule(sigmas)
+    if max_video_history_frames is None:
+        raise ValueError("Streaming rollout requires an explicit history bound")
+    if reference.shape[2] != model.config.num_ref:
+        raise ValueError("Reference latent count differs from model configuration")
+    if reference_vace.shape[2] != reference.shape[2]:
+        raise ValueError("Reference VACE shape mismatch")
+    cache = HistoryCache(max_video_history_frames=max_video_history_frames)
+    base = dict(conditions)
+    with evaluating(model):
+        if model.config.num_ref:
+            ref_conditions = dict(base, vace_context=reference_vace, vace_start=0)
+            commit(model, reference, ref_conditions, 0, cache)
+        expected = model.config.video_blocks(
+            (base["num_frames"]-1)//4+1+model.config.num_ref)
+        for index, item in enumerate(block_inputs):
+            if index >= len(expected):
+                raise ValueError("Too many streaming blocks")
+            start, end = expected[index]
+            noise = item["noise"]
+            if item.get("start", start) != start or noise.shape[2] != end-start:
+                raise ValueError("Streaming block does not match the absolute schedule")
+            local = dict(base, vace_context=item["vace_context"], vace_start=start)
+            current = noise
+            for step in range(len(sigmas)-1):
+                pred = predict_chunk(model, current, sigmas[step], reference, local,
+                                     start, cache, detach_history=True)
+                current = (current.float()+(sigmas[step+1]-sigmas[step])*pred.float()).to(noise.dtype)
+            committed = history_chunk(current, local, start, history_mode)
+            commit(model, committed, local, start, cache)
+            yield {"index": index, "start": start, "end": end,
+                   "chunk": current, "cache_frames": cache.frames,
+                   "cache_retained_frames": cache.retained_frames}
+        if len(expected) and cache.frames != expected[-1][1]:
+            raise ValueError("Streaming input ended before the configured video length")

viddub/causal/train.py

@@ -0,0 +1,6 @@
+"""Backward-compatible generic causal training entry point."""
+from .training.engine import main
+
+
+if __name__ == "__main__":
+    main()

viddub/causal/train_ar.py

@@ -0,0 +1,6 @@
+"""Train the autoregressive causal VidDub parent."""
+from .training.engine import main
+
+
+if __name__ == "__main__":
+    main("ar")

viddub/causal/train_dmd.py

@@ -0,0 +1,6 @@
+"""Refine a few-step VidDub student with distribution matching."""
+from .training.engine import main
+
+
+if __name__ == "__main__":
+    main("dmd")

viddub/causal/train_ode.py

@@ -0,0 +1,6 @@
+"""Distill a few-step ODE VidDub student from frozen trajectories."""
+from .training.engine import main
+
+
+if __name__ == "__main__":
+    main("ode")

viddub/causal/training/__init__.py

@@ -0,0 +1,3 @@
+"""Shared causal training implementation."""
+
+from .stages import ARObjective, DMDObjective, FakeObjective, ODEObjective

viddub/causal/training/ar.py

@@ -0,0 +1,146 @@
+"""Autoregressive causal flow-matching objective."""
+import torch
+
+from ..sampling import evaluating, predict_chunk, schedule
+from .common import ObjectiveBase, flow_noise, spatial_weight, weighted_mse
+
+
+class ARObjective(ObjectiveBase):
+    def __init__(self, model, *, art=None, art_weight=0, art_layer=8,
+                 art_warmup=50, ar_history_mode="gt",
+                 ar_history_prob_start=0, ar_history_prob_end=0,
+                 ar_history_ramp_steps=0, ar_history_steps=4,
+                 ar_history_sigma_max=.5, ar_history_mouth_drop_prob=0,
+                 ar_history_mouth_drop_cross_only=False, **kwargs):
+        super().__init__(model, **kwargs)
+        self.art, self.art_weight = art, art_weight
+        self.art_layer, self.art_warmup = art_layer, art_warmup
+        self.ar_history_mode = ar_history_mode
+        self.ar_history_prob_start = ar_history_prob_start
+        self.ar_history_prob_end = ar_history_prob_end
+        self.ar_history_ramp_steps = ar_history_ramp_steps
+        self.ar_history_steps = ar_history_steps
+        self.ar_history_sigma_max = ar_history_sigma_max
+        self.ar_history_mouth_drop_prob = ar_history_mouth_drop_prob
+        self.ar_history_mouth_drop_cross_only = ar_history_mouth_drop_cross_only
+
+    def _history_probability(self, step):
+        if self.ar_history_ramp_steps <= 0:
+            return self.ar_history_prob_end
+        alpha = min(1., max(0., float(step) / self.ar_history_ramp_steps))
+        return self.ar_history_prob_start + alpha * (
+            self.ar_history_prob_end - self.ar_history_prob_start)
+
+    @torch.no_grad()
+    def _generated_history(self, x0, conditions, end, generator):
+        n_ref = self.model.config.num_ref
+        prefix = x0[:, :, :n_ref].detach()
+        if end <= n_ref:
+            return prefix
+        with evaluating(self.model):
+            for start, stop in self.model.config.video_blocks(end):
+                if self.ar_history_mode == "resample":
+                    u_max = self.ar_history_sigma_max / (
+                        self.shift - self.ar_history_sigma_max * (self.shift - 1))
+                    u = torch.rand((), device=x0.device, generator=generator) * u_max
+                    sigma = self.shift * u / (1 + (self.shift - 1) * u)
+                    clean = x0[:, :, start:stop]
+                    noise = torch.randn(clean.shape, device=clean.device,
+                                        dtype=clean.dtype, generator=generator)
+                    noisy = flow_noise(clean, sigma, noise)
+                    pred = predict_chunk(
+                        self.model, noisy, sigma, prefix, conditions, start,
+                        drop_audio=conditions.get("drop_audio"))
+                    chunk = (noisy.float() - float(sigma) * pred.float()).to(x0.dtype)
+                elif self.ar_history_mode == "rollout":
+                    current = torch.randn(
+                        x0[:, :, start:stop].shape, device=x0.device,
+                        dtype=x0.dtype, generator=generator)
+                    sigmas = schedule(self.ar_history_steps, self.shift, x0.device)
+                    for index in range(len(sigmas) - 1):
+                        pred = predict_chunk(
+                            self.model, current, sigmas[index], prefix,
+                            conditions, start, drop_audio=conditions.get("drop_audio"))
+                        current = (current.float() + (sigmas[index + 1] - sigmas[index])
+                                   * pred.float()).to(x0.dtype)
+                    chunk = current
+                else:
+                    raise ValueError(self.ar_history_mode)
+                prefix = torch.cat([prefix, chunk.detach()], 2)
+        return prefix
+
+    def forward(self, x0, conditions, *, chunk_index=None, sigma=None,
+                noise=None, **_):
+        ranges = self.model.config.video_blocks(x0.shape[2])
+        index = int(torch.randint(len(ranges), ())) if chunk_index is None else chunk_index
+        start, stop = ranges[index]
+        u = torch.rand((), device=x0.device).clamp(1e-4, 1 - 1e-4)
+        sigma = self.shift * u / (1 + (self.shift - 1) * u) if sigma is None else sigma
+        clean = x0[:, :, start:stop]
+        noise = torch.randn_like(clean) if noise is None else noise
+        noisy = flow_noise(clean, sigma, noise)
+        prefix = x0[:, :, :start].detach()
+
+        history_mouth_dropped = False
+        eligible = (start > self.model.config.num_ref and
+                    (not self.ar_history_mouth_drop_cross_only or conditions.get("cross", False)))
+        if self.ar_history_mouth_drop_prob and eligible:
+            generator = torch.Generator(device=x0.device).manual_seed(
+                int(conditions.get("history_seed", 0)) + index * 1009 + 17)
+            history_mouth_dropped = bool(
+                torch.rand((), device=x0.device, generator=generator)
+                < self.ar_history_mouth_drop_prob)
+            if history_mouth_dropped:
+                mask = conditions["vace_context"][:, 32:, :start].amax(
+                    1, keepdim=True).float().clamp(0, 1)
+                mask[:, :, :self.model.config.num_ref] = 0
+                inactive = conditions["vace_context"][:, :x0.shape[1], :start].to(prefix)
+                prefix = prefix * (1 - mask) + inactive * mask
+
+        history_used = False
+        if self.ar_history_mode != "gt" and start > self.model.config.num_ref:
+            generator = torch.Generator(device=x0.device).manual_seed(
+                int(conditions.get("history_seed", 0)))
+            probability = self._history_probability(conditions.get("train_step", 0))
+            history_used = bool(torch.rand((), device=x0.device, generator=generator) < probability)
+            if history_used:
+                prefix = self._generated_history(x0, conditions, start, generator)
+
+        result = predict_chunk(
+            self.model, noisy, sigma, prefix, conditions, start,
+            drop_audio=conditions.get("drop_audio"),
+            gradient_checkpointing=self.gradient_checkpointing,
+            return_hidden_layer=self.art_layer if self.art is not None else None)
+        pred, hidden = result if self.art is not None else (result, None)
+        current_mask = conditions["vace_context"][:, 32:, start:stop].amax(
+            1, keepdim=True).float().clamp(0, 1)
+        weight = ((1 - current_mask) if conditions.get("cross", False)
+                  else spatial_weight(conditions["vace_context"][:, :, start:stop],
+                                      self.mouth_weight))
+        loss = weighted_mse(pred, noise - clean, weight)
+        metrics = {
+            "fm": loss.detach(), "chunk": float(index), "sigma": float(sigma),
+            "history_eligible": float(start > self.model.config.num_ref),
+            "history_generated": float(history_used),
+            "history_probability": float(self._history_probability(
+                conditions.get("train_step", 0))),
+            "cross": float(conditions.get("cross", False)),
+            "history_mouth_dropped": float(history_mouth_dropped),
+        }
+        if self.quality[0] is not None:
+            estimate = noisy.float() - float(sigma) * pred.float()
+            auxiliary, stats = self.quality[0](
+                torch.cat([prefix, estimate], 2), conditions, start, stop, sigma)
+            loss = loss + auxiliary
+            metrics.update({key: value.detach() for key, value in stats.items()})
+            metrics["aux_loss"] = auxiliary.detach()
+        if self.art is not None:
+            target = self.quality[0].art_target(conditions, stop)
+            art_loss = (sum(parameter.sum() * 0 for parameter in self.art.parameters())
+                        if target is None else self.art(
+                            hidden, conditions["pixel_mask"].to(hidden.device), target,
+                            num_ref=self.model.config.num_ref, history_end=start,
+                            warmup=conditions.get("train_step", 0) < self.art_warmup))
+            loss = loss + self.art_weight * art_loss
+            metrics["art"] = art_loss.detach()
+        return loss, metrics

viddub/causal/training/common.py

@@ -0,0 +1,61 @@
+"""Shared math and state used by causal training objectives."""
+import torch
+from torch import nn
+
+
+def spatial_weight(vace_context, mouth_weight):
+    if mouth_weight < 1:
+        raise ValueError("mouth_weight must be >= 1")
+    mask = vace_context[:, 32:].amax(1, keepdim=True).float()
+    return 1 + (mouth_weight - 1) * mask
+
+
+def weighted_mse(a, b, weight):
+    return ((a.float() - b.float()).square() * weight).mean() / weight.mean()
+
+
+def flow_noise(x0, sigma, noise):
+    return (1 - sigma) * x0 + sigma * noise
+
+
+def dmd_gradient(x0, real_x0, fake_x0, active, eps=1e-6):
+    dims = tuple(range(1, x0.ndim))
+    denom = (active * (x0 - real_x0).abs()).sum(dims, keepdim=True)
+    count = active.expand_as(x0).sum(dims, keepdim=True).clamp_min(1)
+    denom = (denom / count).clamp_min(eps)
+    gradient = (fake_x0 - real_x0) / denom
+    if not torch.isfinite(gradient).all():
+        raise FloatingPointError("Non-finite DMD score difference")
+    return gradient * active
+
+
+class ObjectiveBase(nn.Module):
+    """Common generator state; frozen score/quality networks stay unregistered."""
+    def __init__(self, model, *, mouth_weight=3, shift=5, student_steps=4,
+                 use_cache=True, gradient_checkpointing=True, quality=None,
+                 **_ignored):
+        super().__init__()
+        self.model = model
+        self.mouth_weight = mouth_weight
+        self.shift = shift
+        self.student_steps = student_steps
+        self.use_cache = use_cache
+        self.gradient_checkpointing = gradient_checkpointing
+        self.quality = [quality]
+
+
+class ScoreMixin:
+    def _score(self, score, x, timestep, conditions, drop_audio=None, grad=False):
+        if grad and self.dmd_score_offload:
+            raise ValueError("Trainable scores must remain resident through backward")
+        if self.dmd_score_offload:
+            score.to(x.device)
+        try:
+            return score(
+                x, timestep, conditions["context"], conditions["vace_context"],
+                conditions["whisper_hs"], conditions["num_frames"], drop_audio,
+                gradient_checkpointing=grad and self.gradient_checkpointing)
+        finally:
+            if self.dmd_score_offload:
+                score.cpu()
+                torch.cuda.empty_cache()

viddub/causal/training/dmd.py

@@ -0,0 +1,133 @@
+"""Distribution-matching objective for the few-step causal student."""
+import torch
+
+from ..sampling import rollout, schedule
+from .common import (ObjectiveBase, ScoreMixin, dmd_gradient, flow_noise,
+                     spatial_weight, weighted_mse)
+
+
+class DMDObjective(ScoreMixin, ObjectiveBase):
+    def __init__(self, model, *, teacher=None, fake=None, anchor=None,
+                 teacher_cfg=1, matching_min=.02, matching_max=.98,
+                 dmd_rollout_mode="legacy_early_exit",
+                 rollout_quality_weight=1, rollout_quality_cross=False,
+                 dmd_retention_weight=0, dmd_score_offload=False, **kwargs):
+        super().__init__(model, **kwargs)
+        if not 0 < matching_min < matching_max < 1:
+            raise ValueError("Matching range must lie strictly within (0,1)")
+        if dmd_rollout_mode not in ("legacy_early_exit", "final_step"):
+            raise ValueError("Unknown DMD rollout mode")
+        self.scores = [teacher, fake, anchor]
+        self.teacher_cfg = teacher_cfg
+        self.matching_min, self.matching_max = matching_min, matching_max
+        self.dmd_rollout_mode = dmd_rollout_mode
+        self.rollout_quality_weight = rollout_quality_weight
+        self.rollout_quality_cross = rollout_quality_cross
+        self.dmd_retention_weight = dmd_retention_weight
+        self.dmd_score_offload = dmd_score_offload
+
+    def _anchor_rollout(self, anchor, x0, conditions, sigmas, noise):
+        if self.dmd_score_offload:
+            anchor.to(x0.device)
+        try:
+            return rollout(anchor, x0, conditions, sigmas, noise=noise,
+                           use_cache=self.use_cache)
+        finally:
+            if self.dmd_score_offload:
+                anchor.cpu()
+                torch.cuda.empty_cache()
+
+    def forward(self, x0, conditions, *, chunk_index=None,
+                dmd_mode="combined", **_):
+        if dmd_mode not in ("combined", "primary", "quality"):
+            raise ValueError("Unknown DMD backward mode")
+        teacher, fake, anchor = self.scores
+        if teacher is None or fake is None:
+            raise ValueError("DMD requires both real and fake score models")
+        ranges = self.model.config.video_blocks(x0.shape[2])
+        index = int(torch.randint(len(ranges), ())) if chunk_index is None else chunk_index
+        start, stop = ranges[index]
+        exit_step = (int(torch.randint(self.student_steps, ()))
+                     if self.dmd_rollout_mode == "legacy_early_exit" else None)
+        sigmas = schedule(self.student_steps, self.shift, x0.device)
+        rollout_noise = torch.randn_like(x0)
+        generated = rollout(
+            self.model, x0, conditions, sigmas, noise=rollout_noise,
+            use_cache=self.use_cache, grad_chunk=index, grad_step=exit_step,
+            gradient_checkpointing=self.gradient_checkpointing)
+
+        if dmd_mode == "quality":
+            if self.quality[0] is None:
+                raise ValueError("DMD quality-only pass requires quality supervision")
+            if conditions.get("cross", False) and not self.rollout_quality_cross:
+                zero = generated[:, :, start:stop].sum() * 0
+                return zero, {
+                    "rollout_aux_loss": zero.detach(), "chunk": float(index),
+                    "exit_step": float(self.student_steps - 1
+                                       if exit_step is None else exit_step)}
+            auxiliary, stats = self.quality[0](
+                generated[:, :, :stop], conditions, start, stop, 0.)
+            loss = self.rollout_quality_weight * auxiliary
+            metrics = {f"rollout_{key}": value.detach()
+                       for key, value in stats.items()}
+            metrics.update(
+                rollout_aux_loss=auxiliary.detach(), chunk=float(index),
+                exit_step=float(self.student_steps - 1
+                                if exit_step is None else exit_step))
+            return loss, metrics
+
+        with torch.no_grad():
+            sigma = torch.empty((), device=x0.device).uniform_(
+                self.matching_min, self.matching_max)
+            noisy = flow_noise(generated.detach(), sigma, torch.randn_like(generated))
+            timestep = noisy.new_full(
+                (noisy.shape[0], noisy.shape[2]), float(sigma) * 1000)
+            real_velocity = self._score(teacher, noisy, timestep, conditions)
+            if self.teacher_cfg != 1:
+                null = torch.ones(noisy.shape[0], dtype=torch.bool,
+                                  device=noisy.device)
+                null_velocity = self._score(
+                    teacher, noisy, timestep, conditions, null)
+                real_velocity = null_velocity + self.teacher_cfg * (
+                    real_velocity - null_velocity)
+            real_x0 = noisy.float() - sigma * real_velocity.float()
+            fake_x0 = noisy.float() - sigma * self._score(
+                fake, noisy, timestep, conditions).float()
+            active = torch.zeros_like(generated[:, :1], dtype=torch.float32)
+            active[:, :, start:stop] = 1
+            direction = dmd_gradient(
+                generated.float(), real_x0, fake_x0, active)
+            target = generated.detach().float() - direction
+
+        loss = .5 * weighted_mse(
+            generated[:, :, start:stop], target[:, :, start:stop],
+            spatial_weight(conditions["vace_context"][:, :, start:stop],
+                           self.mouth_weight))
+        metrics = {
+            "dmd": loss.detach(),
+            "direction_rms": direction.square().mean().sqrt(),
+            "chunk": float(index),
+            "exit_step": float(self.student_steps - 1
+                               if exit_step is None else exit_step),
+        }
+        if self.dmd_retention_weight:
+            if anchor is None:
+                raise ValueError("DMD retention requires a frozen ODE anchor")
+            with torch.no_grad():
+                anchor_generated = self._anchor_rollout(
+                    anchor, x0, conditions, sigmas, rollout_noise)
+            retention = weighted_mse(
+                generated[:, :, start:stop], anchor_generated[:, :, start:stop],
+                spatial_weight(conditions["vace_context"][:, :, start:stop],
+                               self.mouth_weight))
+            loss = loss + self.dmd_retention_weight * retention
+            metrics["retention"] = retention.detach()
+        if (dmd_mode == "combined" and self.quality[0] is not None and
+                (not conditions.get("cross", False) or self.rollout_quality_cross)):
+            auxiliary, stats = self.quality[0](
+                generated[:, :, :stop], conditions, start, stop, 0.)
+            loss = loss + self.rollout_quality_weight * auxiliary
+            metrics.update({f"rollout_{key}": value.detach()
+                            for key, value in stats.items()})
+            metrics["rollout_aux_loss"] = auxiliary.detach()
+        return loss, metrics

viddub/causal/training/engine.py

@@ -0,0 +1,568 @@
+"""Shared implementation behind the AR, ODE and DMD training entry points."""
+import argparse
+import copy
+import datetime
+import hashlib
+import json
+import random
+import math
+import time
+from pathlib import Path
+
+import torch
+from torch.optim import _functional as optim_functional
+
+from .stages import ARObjective, DMDObjective, FakeObjective, ODEObjective
+from ..runtime import (atomic_json, autocast, checkpoint_paths, export_weights, fingerprint,
+                      first_item, load_components, local_imports, make_dataset, prepare_clip,
+                      read_config, require_causal_metadata)
+from ..sampling import rollout, schedule
+
+
+def optimizer(model, cfg):
+    return torch.optim.AdamW([
+        {"params": model.adapter.parameters(), "lr": cfg["lr"]},
+        {"params": list(model.dit.parameters()) + list(model.vace.parameters()), "lr": cfg["lr_base"]},
+    ], weight_decay=cfg["weight_decay"], betas=(cfg.get('beta1',.9),cfg.get('beta2',.999)), foreach=False)
+
+
+def lr_scheduler(opt, cfg, updates_per_step=1):
+    """Schedule measured in generator updates, including repeated critic steps."""
+    if updates_per_step < 1:
+        raise ValueError("updates_per_step must be positive")
+    def scale(step):
+        step = step / updates_per_step
+        warm = min(1.0, (step + 1) / max(1, cfg["warmup"]))
+        if cfg.get('lr_schedule') == 'cosine' and step >= cfg['warmup']:
+            progress = min(1., (step-cfg['warmup'])/max(1,cfg['max_steps']-cfg['warmup']))
+            return cfg['lr_min_ratio']+(1-cfg['lr_min_ratio'])*0.5*(1+math.cos(math.pi*progress))
+        return warm
+    return torch.optim.lr_scheduler.LambdaLR(opt, scale)
+
+
+def set_trainable(model, enabled):
+    model.requires_grad_(enabled)
+    if not enabled:
+        for p in model.parameters():
+            p.grad = None
+
+
+def move_optimizer_state(opt, device):
+    for state in opt.state.values():
+        for key, value in state.items():
+            if torch.is_tensor(value):
+                state[key] = value.to(device)
+
+
+def restore_optimizer_state_for_step(opt, device):
+    """Restore offloaded Adam state after backward without retaining cache fragments.
+
+    Production clips can leave hundreds of MiB in non-contiguous cached blocks after
+    fake-score backward.  The Adam tensors need contiguous allocations and may OOM
+    despite that memory being reported as reserved-but-free.  Releasing only the
+    allocator cache here does not modify live tensors, gradients, optimizer values,
+    RNG, or update order; it merely makes already-free blocks available to the CUDA
+    allocator before the state transfer.
+    """
+    if torch.device(device).type == "cuda":
+        torch.cuda.empty_cache()
+    move_optimizer_state(opt, device)
+
+
+@torch.no_grad()
+def streaming_offloaded_adamw_step(opt, device, chunk_elements=262144):
+    """Apply the native AdamW update one parameter at a time.
+
+    The optimizer state remains on CPU between parameters.  For each parameter we
+    move bounded chunks of its first/second moments to the parameter device, call
+    PyTorch's own Adam functional with the original parameter-group flags, then
+    immediately move those chunks back to CPU. Every chunk sees the same old step;
+    the persistent step counter advances once per parameter update. This preserves
+    ordinary AdamW state_dict semantics without materializing a full parameter's,
+    let alone the full model's, moments on CUDA.
+    """
+    base = getattr(opt, "optimizer", opt)
+    if not isinstance(base, torch.optim.AdamW):
+        raise TypeError("Streaming offload currently requires torch.optim.AdamW")
+    if chunk_elements < 1:
+        raise ValueError("chunk_elements must be positive")
+    if torch.device(device).type == "cuda":
+        torch.cuda.empty_cache()
+    for group in base.param_groups:
+        if group.get("foreach") not in (False, None) or group.get("fused"):
+            raise ValueError("Streaming offload requires foreach=False and fused=False")
+        beta1, beta2 = group["betas"]
+        for parameter in group["params"]:
+            gradient = parameter.grad
+            if gradient is None:
+                continue
+            if gradient.is_sparse:
+                raise RuntimeError("Adam does not support sparse gradients")
+            state = base.state[parameter]
+            if not state:
+                # Match Adam._init_group, except moments are born on CPU because
+                # this optimizer is explicitly offloaded between updates.
+                state["step"] = torch.tensor(0.0, dtype=torch.float32)
+                state["exp_avg"] = torch.zeros_like(
+                    parameter, device="cpu", memory_format=torch.preserve_format)
+                state["exp_avg_sq"] = torch.zeros_like(
+                    parameter, device="cpu", memory_format=torch.preserve_format)
+                if group["amsgrad"]:
+                    state["max_exp_avg_sq"] = torch.zeros_like(
+                        parameter, device="cpu", memory_format=torch.preserve_format)
+            if not parameter.is_contiguous() or not gradient.is_contiguous():
+                raise ValueError("Streaming AdamW requires contiguous parameters and gradients")
+            parameter_flat = parameter.view(-1)
+            gradient_flat = gradient.view(-1)
+            exp_avg_cpu = state["exp_avg"].view(-1)
+            exp_avg_sq_cpu = state["exp_avg_sq"].view(-1)
+            max_cpu = state.get("max_exp_avg_sq")
+            max_cpu = max_cpu.view(-1) if max_cpu is not None else None
+            old_step = state["step"].clone()
+            next_step = None
+            for start in range(0, parameter.numel(), chunk_elements):
+                end = min(parameter.numel(), start + chunk_elements)
+                exp_avg = exp_avg_cpu[start:end].to(parameter.device)
+                exp_avg_sq = exp_avg_sq_cpu[start:end].to(parameter.device)
+                max_exp_avg_sq = (max_cpu[start:end].to(parameter.device)
+                                  if max_cpu is not None else None)
+                local_step = old_step.clone()
+                optim_functional.adam(
+                    [parameter_flat[start:end]], [gradient_flat[start:end]],
+                    [exp_avg], [exp_avg_sq],
+                    [max_exp_avg_sq] if max_exp_avg_sq is not None else [],
+                    [local_step], foreach=False,
+                    capturable=group["capturable"],
+                    differentiable=group["differentiable"], fused=False,
+                    grad_scale=getattr(base, "grad_scale", None),
+                    found_inf=getattr(base, "found_inf", None),
+                    has_complex=torch.is_complex(parameter),
+                    decoupled_weight_decay=True, amsgrad=group["amsgrad"],
+                    beta1=beta1, beta2=beta2, lr=group["lr"],
+                    weight_decay=group["weight_decay"], eps=group["eps"],
+                    maximize=group["maximize"])
+                exp_avg_cpu[start:end].copy_(exp_avg.cpu())
+                exp_avg_sq_cpu[start:end].copy_(exp_avg_sq.cpu())
+                if max_exp_avg_sq is not None:
+                    max_cpu[start:end].copy_(max_exp_avg_sq.cpu())
+                next_step = local_step
+            if next_step is not None:
+                state["step"].copy_(next_step)
+
+
+def finite_streaming_offloaded_update(accel, objective, opt, cfg):
+    norm = accel.clip_grad_norm_(objective.parameters(), cfg["grad_clip"])
+    valid = accel.reduce(torch.tensor(int(torch.isfinite(norm)), device=accel.device), "sum")
+    if int(valid) != accel.num_processes:
+        opt.zero_grad(set_to_none=True)
+        raise FloatingPointError("Non-finite gradients; stopped without advancing checkpoint")
+    streaming_offloaded_adamw_step(opt, accel.device)
+    opt.zero_grad(set_to_none=True)
+    return float(norm)
+
+
+def finite_update(accel, objective, opt, cfg):
+    norm = accel.clip_grad_norm_(objective.parameters(), cfg["grad_clip"])
+    valid = accel.reduce(torch.tensor(int(torch.isfinite(norm)), device=accel.device), "sum")
+    if int(valid) != accel.num_processes:
+        opt.zero_grad(set_to_none=True)
+        raise FloatingPointError("Non-finite gradients; stopped without advancing checkpoint")
+    opt.step()
+    opt.zero_grad(set_to_none=True)
+    return float(norm)
+
+
+def check_loss(accel, loss):
+    finite = accel.reduce(torch.tensor(int(torch.isfinite(loss)), device=accel.device), "sum")
+    if int(finite) != accel.num_processes:
+        raise FloatingPointError("Non-finite loss on at least one rank")
+
+
+def main(expected_stage=None):
+    p = argparse.ArgumentParser(description=__doc__)
+    p.set_defaults(teacher_dir=None, teacher_step=None, fake_dir=None,
+                   fake_step=None, trajectories=None)
+    p.add_argument("--config", required=True)
+    p.add_argument("--init-dir", required=True)
+    p.add_argument("--init-step", required=True, type=int)
+    p.add_argument("--output", required=True)
+    if expected_stage in (None, "ode"):
+        p.add_argument(
+            "--trajectories",
+            help="Trajectory directory produced by viddub.causal.make_trajectories",
+        )
+    if expected_stage in (None, "dmd"):
+        p.add_argument("--teacher-dir", help="Original bidirectional VidDub checkpoint")
+        p.add_argument("--teacher-step", type=int)
+        p.add_argument("--fake-dir", help="Pretrained fake-score checkpoint")
+        p.add_argument("--fake-step", type=int)
+    p.add_argument("--resume", help="A complete state_stepNNNNNN directory")
+    p.add_argument("--stop-after-step", type=int,
+                   help="Operational stop after a saved step; does not alter the training config")
+    p.add_argument("--check-only", action="store_true", help="Validate inputs, metadata and data; do not allocate GPUs")
+    a = p.parse_args()
+    cfg = read_config(a.config)
+    if expected_stage is not None and cfg["stage"] != expected_stage:
+        p.error(f"This entry point requires stage={expected_stage}, got {cfg['stage']}")
+    if cfg['stage'] == 'cd':
+        p.error('Use python -m viddub.causal.cf.train for consistency training')
+    if a.stop_after_step is not None and not 0 < a.stop_after_step <= cfg['max_steps']:
+        p.error('--stop-after-step must lie within 1..max_steps')
+    run_until = cfg['max_steps'] if a.stop_after_step is None else a.stop_after_step
+    paths = checkpoint_paths(a.init_dir, a.init_step)
+    teacher_paths = None
+    run_provenance = {"student_init": fingerprint(paths)}
+    if cfg["stage"] == "ode":
+        if cfg["causal"]:
+            require_causal_metadata(a.init_dir, a.init_step, cfg, ("ar",))
+        if not a.trajectories:
+            p.error("ODE requires --trajectories")
+    if cfg["stage"] == "dmd":
+        require_causal_metadata(a.init_dir, a.init_step, cfg, ("ode", "dmd", "cd"))
+        teacher_paths = checkpoint_paths(a.teacher_dir, a.teacher_step)
+        run_provenance["teacher"] = fingerprint(teacher_paths)
+        if bool(a.fake_dir) != (a.fake_step is not None):
+            p.error("Provide both --fake-dir and --fake-step")
+        fake_paths = checkpoint_paths(a.fake_dir, a.fake_step) if a.fake_dir else None
+        if fake_paths is not None:
+            run_provenance['fake_init'] = fingerprint(fake_paths)
+    if a.resume:
+        if Path(a.resume).resolve().parent != Path(a.output).resolve():
+            raise ValueError("Resume state must belong to this output directory")
+        resume_meta = json.loads((Path(a.resume) / "complete.json").read_text())
+        if read_config(resume_meta["config"]) != cfg:
+            raise ValueError("Resume requires exactly the saved config")
+    else:
+        resume_meta = None
+    if not a.resume and Path(a.output).exists() and any(Path(a.output).iterdir()):
+        raise ValueError("Output is not empty; choose a new output or explicitly --resume")
+    if cfg["stage"] == "ode":
+        from ..trajectories import TrajectoryDataset
+        dataset = TrajectoryDataset(a.trajectories, cfg)
+        trajectory_causal = dataset.manifest.get(
+            "teacher_architecture", {}
+        ).get("causal", dataset.manifest.get("config", {}).get("causal", True))
+        if trajectory_causal != cfg["causal"]:
+            raise ValueError("Trajectory/student causal architecture mismatch")
+        for name in paths:
+            if dataset.manifest["teacher"][name]["sha256"] != run_provenance["student_init"][name]["sha256"]:
+                raise ValueError(f"Trajectory teacher differs from student initialization: {name}")
+        run_provenance["trajectories"] = str(Path(a.trajectories).resolve())
+        run_provenance["trajectory_manifest"] = fingerprint({"manifest": dataset.manifest_path})["manifest"]
+        rows = None
+    else:
+        dataset, rows = make_dataset(cfg)
+    if cfg["stage"] == "dmd" and cfg["p_cross"] > 0 and len(dataset) < 2:
+        raise ValueError("Cross-audio DMD requires at least two training clips")
+    if rows is not None:
+        row_json = json.dumps(rows, sort_keys=True).encode()
+        run_provenance["dataset_rows_sha256"] = hashlib.sha256(row_json).hexdigest()
+    if resume_meta and resume_meta["provenance"] != run_provenance:
+        raise ValueError("Resume checkpoint/data provenance differs")
+    if a.check_only:
+        print(json.dumps({"config": cfg, "checkpoint_files": {k: str(v) for k, v in paths.items()},
+                          "samples": len(dataset), "status": "inputs_validated_no_gpu_execution"}, indent=2))
+        return
+    local_imports()
+    import accelerate
+    from safetensors.torch import load_file
+    from viddub.full_context.runtime import ZeroLatents
+    accel = accelerate.Accelerator(
+        mixed_precision="no", gradient_accumulation_steps=cfg["grad_accum"],
+        step_scheduler_with_optimizer=False,
+        kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=False, gradient_as_bucket_view=True),
+                         accelerate.InitProcessGroupKwargs(timeout=datetime.timedelta(hours=2))])
+    if accel.device.type != "cuda":
+        raise RuntimeError(
+            "Production training needs CUDA; run `python -m unittest discover -s tests` "
+            "for CPU-only contract verification"
+        )
+    if "DEEPSPEED" in str(accel.distributed_type) or "FSDP" in str(accel.distributed_type):
+        raise ValueError("This first implementation supports plain single GPU/DDP only")
+    random.seed(cfg["seed"] + accel.process_index)
+    torch.manual_seed(cfg["seed"] + accel.process_index)
+    pipe, whisper, student = load_components(paths, cfg, accel.device)
+    student.requires_grad_(True).train()
+    teacher = fake = anchor = fake_objective = fake_opt = fake_sched = None
+    if cfg["stage"] == "dmd":
+        if cfg['dmd_retention_weight']:
+            anchor = copy.deepcopy(student).requires_grad_(False).eval()
+        teacher = copy.deepcopy(student)
+        teacher.causal = False
+        for name, path in teacher_paths.items():
+            getattr(teacher, name).load_state_dict(load_file(str(path)), strict=True)
+        teacher.requires_grad_(False).eval()
+        fake = copy.deepcopy(teacher)
+        if fake_paths is not None:
+            for name, path in fake_paths.items():
+                getattr(fake, name).load_state_dict(load_file(str(path)), strict=True)
+        fake.requires_grad_(True).train()
+    kwargs = dict(mouth_weight=cfg["mouth_weight"], shift=cfg["shift"], student_steps=cfg["student_steps"],
+                  teacher_cfg=cfg["teacher_cfg"], use_cache=cfg["use_cache"],
+                  gradient_checkpointing=cfg["gradient_checkpointing"],
+                  matching_min=cfg["matching_min"], matching_max=cfg["matching_max"])
+    kwargs['dmd_rollout_mode'] = cfg['dmd_rollout_mode']
+    kwargs['ode_gt_flow_weight'] = cfg['ode_gt_flow_weight']
+    kwargs['rollout_quality_weight'] = cfg['rollout_quality_weight']
+    kwargs['rollout_quality_cross'] = cfg['rollout_quality_cross']
+    kwargs['dmd_retention_weight'] = cfg['dmd_retention_weight']
+    kwargs['dmd_score_offload'] = cfg['dmd_score_offload']
+    for key in ['ar_history_mode','ar_history_prob_start','ar_history_prob_end',
+                'ar_history_ramp_steps','ar_history_steps','ar_history_sigma_max',
+                'ar_history_mouth_drop_prob','ar_history_mouth_drop_cross_only']:
+        kwargs[key]=cfg[key]
+    quality = None
+    if any(cfg.get(k,0) for k in ['pixel_weight','lpips_weight','sync_weight','keep_weight','art_weight']):
+        from ..quality import GTQuality
+        quality = GTQuality(pipe, cfg, accel.device)
+    art = None
+    if cfg.get('art_weight',0):
+        if cfg['stage']!='ar':raise ValueError('Art integration is initially AR-only')
+        from ..art_alignment import CausalArtAlignment
+        art=CausalArtAlignment(student.dit.dim).to(accel.device)
+    objective_class={"ar":ARObjective,"ode":ODEObjective,"dmd":DMDObjective}[cfg["stage"]]
+    if cfg['stage']=='dmd' and cfg['sampler']=='cf_renoise':
+        from ..cf.dmd import CFAsymmetricDMD
+        objective_class=CFAsymmetricDMD
+    objective = objective_class(student, teacher=teacher, fake=fake, anchor=anchor,
+                          quality=quality,
+                          art=art,art_weight=cfg['art_weight'],art_layer=cfg['art_layer'],
+                          art_warmup=cfg['art_warmup'],**kwargs)
+    opt = optimizer(student, cfg)
+    if art is not None:opt.add_param_group({'params':art.parameters(),'lr':cfg['art_lr']})
+    sched = lr_scheduler(opt, cfg)
+    loader_kwargs = dict(batch_size=1, shuffle=True, collate_fn=first_item, num_workers=cfg["workers"])
+    loader_kwargs["generator"] = torch.Generator().manual_seed(cfg["seed"])
+    if cfg["workers"]:
+        loader_kwargs.update(persistent_workers=True, prefetch_factor=2)
+    loader = torch.utils.data.DataLoader(dataset, **loader_kwargs)
+    objective, opt, loader = accel.prepare(objective, opt, loader)
+    accel.register_for_checkpointing(sched)
+    if fake is not None:
+        # The training loop owns fake residency through forward+backward.
+        # Self-offloading inside Objective.fake would move checkpointed weights
+        # to CPU before backward recomputation.
+        fake_kwargs = dict(kwargs, dmd_score_offload=False)
+        fake_objective = FakeObjective(fake, **fake_kwargs)
+        fake_opt = optimizer(fake, dict(cfg, lr=cfg['fake_lr'], lr_base=cfg['fake_lr_base']))
+        fake_sched = lr_scheduler(fake_opt, cfg, updates_per_step=cfg["fake_updates"])
+        fake_objective, fake_opt = accel.prepare(fake_objective, fake_opt)
+        accel.register_for_checkpointing(fake_sched)
+        if cfg['dmd_score_offload']:
+            teacher.cpu()
+            if anchor is not None:anchor.cpu()
+            fake.cpu()
+            torch.cuda.empty_cache()
+    provenance = run_provenance
+    out = Path(a.output)
+    if accel.is_main_process:
+        out.mkdir(parents=True, exist_ok=True)
+        atomic_json(out / "run.json", {"config": cfg, "provenance": provenance,
+                                      "world_size": accel.num_processes,
+                                      "note": "DDP shards samples; all models are replicated per GPU."})
+    accel.wait_for_everyone()
+    provenance = json.loads((out / "run.json").read_text())["provenance"]
+    step, epoch, consumed = 0, 0, 0
+    if a.resume:
+        if resume_meta["world_size"] != accel.num_processes:
+            raise ValueError("Resume requires the same world size")
+        accel.load_state(a.resume)
+        if fake is not None and cfg['dmd_score_offload']:
+            move_optimizer_state(fake_opt, 'cpu')
+            fake.cpu()
+            teacher.cpu()
+            if anchor is not None:anchor.cpu()
+            torch.cuda.empty_cache()
+        step, epoch, consumed = (resume_meta[k] for k in ("step", "epoch", "consumed_batches"))
+    zero_latents = ZeroLatents(pipe)
+    # Auxiliary network initialization must not change the data/noise stream.
+    if not a.resume:
+        random.seed(cfg['seed']+accel.process_index)
+        torch.manual_seed(cfg['seed']+accel.process_index)
+    start_time = time.monotonic()
+    row_index = {r["clip_id"]: i for i, r in enumerate(rows)} if rows is not None else {}
+    metric_sums = {}
+    metric_count = 0
+
+    def save(epoch, consumed):
+        state = out / f"state_step{step:06d}"
+        # All ranks must finish the existence check BEFORE any rank calls
+        # save_state (which creates the directory). Otherwise a faster rank
+        # makes a slower rank mistake this very save for an overwrite.
+        exists = accel.reduce(torch.tensor(int(state.exists()), device=accel.device), "sum")
+        if int(exists):
+            raise FileExistsError(f"Refusing to overwrite state {state}")
+        full_state = not cfg.get('save_state_every') or step % cfg['save_state_every'] == 0 or step == cfg['max_steps']
+        if full_state:accel.save_state(str(state))
+        if accel.is_main_process:
+            export_weights(student, out, step, cfg, provenance)
+            if full_state:atomic_json(state / "complete.json", {"step": step, "epoch": epoch,
+                        "consumed_batches": consumed, "config": cfg, "provenance": provenance,
+                        "world_size": accel.num_processes})
+        accel.wait_for_everyone()
+
+    while step < run_until:
+        if hasattr(loader, "set_epoch"):
+            loader.set_epoch(epoch)
+        # Seed permutation by epoch, independent of model/Whisper RNG use.
+        if getattr(loader, "generator", None) is not None:
+            loader.generator.manual_seed(cfg["seed"] + epoch)
+        active_loader = accel.skip_first_batches(loader, consumed) if consumed else loader
+        for item in active_loader:
+            consumed += 1
+            record = None
+            if cfg["stage"] == "ode":
+                dtype = torch.bfloat16 if cfg["precision"] == "bf16" else torch.float32
+                x0 = item["shape_reference"].to(accel.device, dtype)
+                c = {}
+                for k, v in item["conditions"].items():
+                    if torch.is_tensor(v):
+                        c[k] = v.to(accel.device, dtype) if k in (
+                            'context', 'vace_context', 'whisper_hs') else v
+                    else:
+                        c[k] = v
+                record = item["record"]
+            else:
+                if cfg["stage"] in ("ar", "dmd") and random.random() < cfg["p_cross"]:
+                    # Conditions stay on the source video; swap only audio.
+                    other = dataset[random.randrange(len(dataset))]
+                    if other["clip_id"] == item["clip_id"]:
+                        source_index = row_index[item["clip_id"]]
+                        other = dataset[(source_index + 1) % len(dataset)]
+                    item = dict(item, audio=other["audio"], cross=True)
+                x0, c = prepare_clip(item, pipe, whisper, zero_latents, cfg)
+                c['cross'] = bool(item.get('cross', False))
+            if cfg["stage"] == "ar":
+                c['train_step']=step
+                c['history_seed']=cfg['seed']+consumed*100003+accel.process_index
+                c["drop_audio"] = torch.rand(x0.shape[0], device=x0.device) < cfg["p_drop_audio"]
+                if c['cross']:
+                    c["drop_audio"].zero_()
+            if fake is not None:
+                set_trainable(student, False)
+                if cfg['dmd_score_offload']:
+                    fake.to(accel.device)
+                set_trainable(fake, True)
+                fake.train()
+                for _ in range(cfg["fake_updates"]):
+                    # Keep the (large) Adam state on CPU while rolling out the
+                    # student sample. After the first fake update those states
+                    # otherwise overlap the causal rollout peak by only tens of
+                    # MiB on 96-GB cards. This changes only residency timing:
+                    # fake weights, samples, RNG, update count and objectives are
+                    # unchanged.
+                    with torch.no_grad(), autocast(accel.device, cfg["precision"]):
+                        if cfg['sampler']=='cf_renoise':
+                            from ..cf.dmd import shared_exit
+                            from ..cf.sampling import cf_rollout
+                            generated=cf_rollout(student,x0,c,steps=cfg['student_steps'],shift=cfg['shift'],
+                                use_cache=cfg['use_cache'],grad_exit=(None if cfg['dmd_rollout_mode']=='final_step'
+                                else shared_exit(cfg['student_steps'],x0.device)))
+                        else:
+                            generated = rollout(student, x0, c, schedule(cfg["student_steps"], cfg["shift"], x0.device),
+                                                use_cache=cfg["use_cache"])
+                    with autocast(accel.device, cfg["precision"]):
+                        fake_loss, _ = fake_objective(x0, c, generated=generated)
+                    check_loss(accel, fake_loss)
+                    accel.backward(fake_loss)
+                    if cfg['dmd_score_offload']:
+                        finite_streaming_offloaded_update(
+                            accel, fake_objective, fake_opt, cfg)
+                    else:
+                        finite_update(accel, fake_objective, fake_opt, cfg)
+                    fake_sched.step()
+                    del generated
+                set_trainable(fake, False)
+                fake.eval()
+                if cfg['dmd_score_offload']:
+                    fake.cpu()
+                    torch.cuda.empty_cache()
+                set_trainable(student, True)
+            with accel.accumulate(objective):
+                if cfg['stage'] == 'ar' and cfg['ar_blocks_per_sample'] != 1:
+                    ranges = student.config.video_blocks(x0.shape[2])
+                    count = len(ranges) if cfg['ar_blocks_per_sample'] == 0 else cfg['ar_blocks_per_sample']
+                    if count > len(ranges):
+                        raise ValueError(f'ar_blocks_per_sample={count} exceeds {len(ranges)} video blocks')
+                    # Without replacement: every selected block provides a distinct
+                    # temporal/audio location. Zero means the official-style all-block path.
+                    indices = (list(range(len(ranges))) if count == len(ranges)
+                               else torch.randperm(len(ranges))[:count].tolist())
+                    metrics = {}
+                    loss_value = x0.new_zeros((), dtype=torch.float32)
+                    for chunk_index in indices:
+                        with autocast(accel.device, cfg["precision"]):
+                            chunk_loss, chunk_metrics = objective(x0, c, chunk_index=chunk_index)
+                        check_loss(accel, chunk_loss)
+                        accel.backward(chunk_loss / count)
+                        loss_value = loss_value + chunk_loss.detach().float() / count
+                        for key, value in chunk_metrics.items():
+                            value = torch.as_tensor(value, device=x0.device, dtype=torch.float32).detach()
+                            metrics[key] = metrics.get(key, 0.0) + value / count
+                    loss = loss_value
+                    metrics['blocks_per_sample'] = x0.new_tensor(float(count))
+                    metrics['block_coverage_fraction'] = x0.new_tensor(float(count / len(ranges)))
+                else:
+                    split_quality = (cfg['stage'] == 'dmd' and cfg['dmd_split_quality_backward'] and
+                                     any(cfg[k] for k in ['pixel_weight','lpips_weight','sync_weight','keep_weight']))
+                    if split_quality:
+                        cpu_rng = torch.random.get_rng_state()
+                        cuda_rng = torch.cuda.get_rng_state(accel.device)
+                        with autocast(accel.device, cfg["precision"]):
+                            primary, metrics = objective(x0, c, record=record, dmd_mode='primary')
+                        post_cpu_rng = torch.random.get_rng_state()
+                        post_cuda_rng = torch.cuda.get_rng_state(accel.device)
+                        check_loss(accel, primary)
+                        accel.backward(primary)
+                        torch.random.set_rng_state(cpu_rng)
+                        torch.cuda.set_rng_state(cuda_rng, accel.device)
+                        with autocast(accel.device, cfg["precision"]):
+                            auxiliary, aux_metrics = objective(x0, c, record=record, dmd_mode='quality')
+                        check_loss(accel, auxiliary)
+                        accel.backward(auxiliary)
+                        torch.random.set_rng_state(post_cpu_rng)
+                        torch.cuda.set_rng_state(post_cuda_rng, accel.device)
+                        loss = primary.detach().float() + auxiliary.detach().float()
+                        metrics.update(aux_metrics)
+                        metrics['split_quality_backward'] = x0.new_tensor(1.)
+                    else:
+                        with autocast(accel.device, cfg["precision"]):
+                            loss, metrics = objective(x0, c, record=record)
+                        check_loss(accel, loss)
+                        accel.backward(loss)
+                if accel.sync_gradients:
+                    norm = finite_update(accel, objective, opt, cfg)
+                    sched.step()
+            metric_count += 1
+            for key, value in metrics.items():
+                metric_sums[key] = metric_sums.get(key, 0.0) + float(value)
+            if not accel.sync_gradients:
+                continue
+            step += 1
+            if step % cfg["log_every"] == 0:
+                value = float(accel.reduce(loss.detach(), "mean"))
+                count = accel.reduce(torch.tensor(float(metric_count), device=accel.device), "sum")
+                averaged = {}
+                for key, total in metric_sums.items():
+                    reduced = accel.reduce(torch.tensor(total, device=accel.device), "sum")
+                    averaged[key] = float(reduced / count)
+                log = {"step": step, "loss": value, "grad_norm": norm,
+                       "elapsed_seconds": time.monotonic() - start_time,
+                       "metric_window_steps": cfg["log_every"], **averaged}
+                if fake is not None:
+                    log["fake_loss"] = float(fake_loss.detach())
+                accel.print(json.dumps(log), flush=True)
+                if accel.is_main_process:
+                    with (out / "train.jsonl").open("a") as f:
+                        f.write(json.dumps(log) + "\n")
+                metric_sums.clear();metric_count=0
+            if step % cfg["save_every"] == 0 or step in cfg.get('save_steps',[]) or step == cfg["max_steps"]:
+                save(epoch, consumed)
+            if step == run_until:
+                break
+        epoch, consumed = epoch + 1, 0
+    accel.end_training()
+
+
+if __name__ == "__main__":
+    main()

viddub/causal/training/fake.py

@@ -0,0 +1,34 @@
+"""Fake-score flow-matching objective used by DMD."""
+import torch
+
+from .common import (ObjectiveBase, ScoreMixin, flow_noise, spatial_weight,
+                     weighted_mse)
+
+
+class FakeObjective(ScoreMixin, ObjectiveBase):
+    def __init__(self, model, *, matching_min=.02, matching_max=.98,
+                 dmd_score_offload=False, **kwargs):
+        super().__init__(model, **kwargs)
+        if not 0 < matching_min < matching_max < 1:
+            raise ValueError("Matching range must lie strictly within (0,1)")
+        self.matching_min, self.matching_max = matching_min, matching_max
+        self.dmd_score_offload = dmd_score_offload
+
+    def forward(self, x0, conditions, *, generated=None, **_):
+        if generated is None:
+            raise ValueError("Fake score training requires a student rollout")
+        generated = generated.detach()
+        sigma = torch.empty((), device=generated.device).uniform_(
+            self.matching_min, self.matching_max)
+        noise = torch.randn_like(generated)
+        noisy = flow_noise(generated, sigma, noise)
+        n_ref = self.model.config.num_ref
+        timestep = generated.new_full(
+            (generated.shape[0], generated.shape[2]), float(sigma) * 1000)
+        pred = self._score(
+            self.model, noisy, timestep, conditions, grad=True)
+        loss = weighted_mse(
+            pred[:, :, n_ref:], (noise - generated)[:, :, n_ref:],
+            spatial_weight(conditions["vace_context"][:, :, n_ref:],
+                           self.mouth_weight))
+        return loss, {"fake_flow": loss.detach()}

viddub/causal/training/ode.py

@@ -0,0 +1,72 @@
+"""Trajectory-regression objective for few-step ODE distillation."""
+import torch
+
+from ..sampling import predict_chunk, rollout, schedule
+from .common import ObjectiveBase, flow_noise, spatial_weight, weighted_mse
+
+
+class ODEObjective(ObjectiveBase):
+    def __init__(self, model, *, ode_gt_flow_weight=0,
+                 rollout_quality_weight=1, rollout_quality_cross=False,
+                 **kwargs):
+        super().__init__(model, **kwargs)
+        self.ode_gt_flow_weight = ode_gt_flow_weight
+        self.rollout_quality_weight = rollout_quality_weight
+        self.rollout_quality_cross = rollout_quality_cross
+
+    def forward(self, x0, conditions, *, record=None, **_):
+        if record is None:
+            raise ValueError("ODE stage requires a trajectory record")
+        start, stop, sigma = record["start"], record["end"], record["sigma"]
+        noisy = record["noisy"].to(x0)
+        pred = predict_chunk(
+            self.model, noisy, sigma, record["prefix"].to(x0), conditions, start,
+            gradient_checkpointing=self.gradient_checkpointing)
+        estimate = noisy.float() - float(sigma) * pred.float()
+        loss = weighted_mse(
+            estimate, record["target"].to(x0),
+            spatial_weight(conditions["vace_context"][:, :, start:stop],
+                           self.mouth_weight))
+        metrics = {"ode_x0": loss.detach()}
+
+        gt = record.get("gt_x0")
+        if self.ode_gt_flow_weight:
+            if gt is None:
+                raise ValueError("ODE GT flow requires trajectories containing gt_latents")
+            if not conditions.get("cross", False):
+                gt = gt.to(x0)
+                clean = gt[:, :, start:stop]
+                u = torch.rand((), device=x0.device).clamp(1e-4, 1 - 1e-4)
+                gt_sigma = self.shift * u / (1 + (self.shift - 1) * u)
+                noise = torch.randn_like(clean)
+                gt_noisy = flow_noise(clean, gt_sigma, noise)
+                gt_pred = predict_chunk(
+                    self.model, gt_noisy, gt_sigma, gt[:, :, :start].detach(),
+                    conditions, start,
+                    gradient_checkpointing=self.gradient_checkpointing)
+                gt_flow = weighted_mse(
+                    gt_pred, noise - clean,
+                    spatial_weight(conditions["vace_context"][:, :, start:stop],
+                                   self.mouth_weight))
+                loss = loss + self.ode_gt_flow_weight * gt_flow
+                metrics["ode_gt_flow"] = gt_flow.detach()
+            else:
+                metrics["ode_gt_flow"] = loss.detach() * 0
+
+        if (self.quality[0] is not None and
+                (not conditions.get("cross", False) or self.rollout_quality_cross)):
+            ranges = self.model.config.video_blocks(x0.shape[2])
+            chunk_index = next(
+                index for index, pair in enumerate(ranges) if pair == (start, stop))
+            generated = rollout(
+                self.model, x0, conditions,
+                schedule(self.student_steps, self.shift, x0.device),
+                use_cache=self.use_cache, grad_chunk=chunk_index, grad_step=None,
+                gradient_checkpointing=self.gradient_checkpointing)
+            auxiliary, stats = self.quality[0](
+                generated[:, :, :stop], conditions, start, stop, 0.)
+            loss = loss + self.rollout_quality_weight * auxiliary
+            metrics.update({f"rollout_{key}": value.detach()
+                            for key, value in stats.items()})
+            metrics["rollout_aux_loss"] = auxiliary.detach()
+        return loss, metrics

viddub/causal/training/stages.py

@@ -0,0 +1,5 @@
+"""Stable imports for the public training stages."""
+from .ar import ARObjective
+from .dmd import DMDObjective
+from .fake import FakeObjective
+from .ode import ODEObjective

viddub/causal/trajectories.py

@@ -0,0 +1,66 @@
+"""Validated on-disk causal teacher trajectories (one clip per .pt shard)."""
+import json
+from pathlib import Path
+
+import torch
+
+from .sampling import schedule
+
+
+class TrajectoryDataset(torch.utils.data.Dataset):
+    def __init__(self, directory, cfg):
+        source = Path(directory).resolve()
+        self.manifest_path = source if source.is_file() else source / 'manifest.json'
+        self.directory = self.manifest_path.parent
+        self.cfg = cfg
+        self.manifest = json.loads(self.manifest_path.read_text())
+        if self.manifest.get("schema") != 1 or not self.manifest.get("complete"):
+            raise ValueError("Trajectory generation has not completed")
+        for key in ("chunk_size", "num_ref", "fps", "num_frames", "shift"):
+            if self.manifest["config"][key] != cfg[key]:
+                raise ValueError(f"Trajectory/config mismatch for {key}")
+        if self.manifest["student_steps"] != cfg["student_steps"]:
+            raise ValueError("Trajectory student step count mismatch")
+        if self.manifest["teacher_cfg"] != cfg["teacher_cfg"]:
+            raise ValueError("Trajectory guidance mismatch")
+        self.records = self.manifest["records"]
+        if not self.records:
+            raise ValueError("No trajectories")
+        self.indices = self.manifest["boundary_indices"][:-1]
+        self.sigmas = self.manifest["teacher_sigmas"]
+        expected = schedule(cfg["student_steps"], cfg["shift"])
+        actual = torch.tensor([self.sigmas[i] for i in self.manifest["boundary_indices"]])
+        if not torch.allclose(actual, expected, atol=1e-6, rtol=1e-6):
+            raise ValueError("Teacher trajectory boundaries do not match student schedule")
+        for record in self.records:
+            path = (self.directory / record["file"]).resolve()
+            if not path.is_relative_to(self.directory) or not path.is_file():
+                raise ValueError(f"Missing/unsafe trajectory path: {path}")
+            if "bytes" in record and path.stat().st_size != record["bytes"]:
+                raise ValueError(f"Trajectory size mismatch: {path}")
+
+    def __len__(self):
+        return len(self.records)
+
+    def __getitem__(self, index):
+        record = self.records[index]
+        data = torch.load(self.directory / record["file"], map_location="cpu", weights_only=True)
+        if any(self.cfg.get(k, 0) for k in ('ode_gt_flow_weight','pixel_weight','lpips_weight',
+                                            'sync_weight','keep_weight')):
+            if data.get('gt_latents') is None:
+                raise ValueError('This ODE variant requires trajectories with gt_latents')
+            if data.get('conditions', {}).get('gt_frames') is None:
+                raise ValueError('This ODE variant requires trajectories with pixel GT')
+        chunks = data["chunks"]
+        chunk = chunks[int(torch.randint(len(chunks), ()))]
+        step = int(torch.randint(len(self.indices), ()))
+        i = self.indices[step]
+        a, b = chunk["start"], chunk["end"]
+        conditions = dict(data["conditions"])
+        conditions["cross"] = data.get("mode", record.get("mode", "same")) == "cross"
+        return {"shape_reference": data["shape_reference"], "conditions": conditions,
+                "record": {"start": a, "end": b, "sigma": self.sigmas[i],
+                           "noisy": chunk["states"][i], "target": chunk["states"][-1],
+                           "prefix": data["generated"][:, :, :a],
+                           "gt_x0": data.get("gt_latents")},
+                "clip_id": data["clip_id"]}

viddub/causal/verify_checkpoint.py

@@ -0,0 +1,101 @@
+"""Verify real Wan-Dub weights before spending GPU time on causal training.
+
+Checks original bidirectional parity, future-video isolation, cache/no-cache
+parity, and one AR backward pass. Does not step an optimizer or save weights.
+"""
+import argparse
+from pathlib import Path
+import time
+
+import torch
+
+from .runtime import (atomic_json, autocast, checkpoint_paths, load_components,
+                      local_imports, make_dataset, prepare_clip, read_config)
+from .model import HistoryCache
+from .training.stages import ARObjective
+from .sampling import commit, predict_chunk
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--config", required=True)
+    p.add_argument("--checkpoint-dir", required=True)
+    p.add_argument("--step", type=int, required=True)
+    p.add_argument("--output", required=True, help="New JSON report path")
+    p.add_argument("--clip-index", type=int, default=0)
+    a = p.parse_args()
+    cfg = read_config(a.config)
+    paths = checkpoint_paths(a.checkpoint_dir, a.step)
+    if Path(a.output).exists():
+        raise FileExistsError(a.output)
+    if not torch.cuda.is_available():
+        raise RuntimeError("Real checkpoint verification requires CUDA")
+    local_imports()
+    from viddub.full_context.runtime import ZeroLatents, predict
+    data, _ = make_dataset(cfg)
+    device = torch.device("cuda")
+    pipe, whisper, model = load_components(paths, cfg, device)
+    model.eval().requires_grad_(False)
+    x0, c = prepare_clip(data[a.clip_index], pipe, whisper, ZeroLatents(pipe), cfg)
+    started = time.monotonic()
+    # Native reference uses the project's own adapter ordering and attention backend.
+    torch.manual_seed(cfg["seed"])
+    xt = torch.randn_like(x0)
+    tol = 0.04 if cfg["precision"] == "bf16" else 0.0003
+    report = {"status": "running", "precision": cfg["precision"], "clip_id": data[a.clip_index]["clip_id"],
+              "tolerance": {"atol": tol, "rtol": tol}, "checks": {}}
+    try:
+        with torch.no_grad(), autocast(device, cfg["precision"]):
+            tok = model.adapter.tokens(c["whisper_hs"], cfg["num_frames"], cfg["num_ref"],
+                                       torch.tensor([False], device=device))
+            native = predict(pipe, model.adapter, xt, xt.new_tensor([500]), c["context"], c["vace_context"], tok)
+            model.causal = False
+            actual = model(xt, 500, c["context"], c["vace_context"], c["whisper_hs"], c["num_frames"])
+            report["checks"]["bidirectional_max_abs"] = float((actual - native).abs().max())
+            torch.testing.assert_close(actual, native, atol=tol, rtol=tol)
+            model.causal = True
+            full = model(xt, 500, c["context"], c["vace_context"], c["whisper_hs"], c["num_frames"])
+            end = model.config.video_blocks(xt.shape[2])[0][1]
+            changed, vc = xt.clone(), c["vace_context"].clone()
+            changed[:, :, end:] += 20
+            vc[:, :, end:] -= 20
+            future = model(changed, 500, c["context"], vc, c["whisper_hs"], c["num_frames"])
+            report["checks"]["future_max_abs"] = float((full[:, :, :end] - future[:, :, :end]).abs().max())
+            torch.testing.assert_close(full[:, :, :end], future[:, :, :end], atol=tol, rtol=tol)
+            cache = HistoryCache()
+            n_ref = cfg["num_ref"]
+            if n_ref:
+                commit(model, x0[:, :, :n_ref], c, 0, cache)
+            errors = []
+            for start, end in model.config.video_blocks(x0.shape[2]):
+                expected = predict_chunk(model, xt[:, :, start:end], 0.5, x0[:, :, :start], c, start)
+                actual = predict_chunk(model, xt[:, :, start:end], 0.5, x0[:, :, :start], c, start, cache)
+                errors.append(float((actual - expected).abs().max()))
+                torch.testing.assert_close(actual, expected, atol=tol, rtol=tol)
+                commit(model, x0[:, :, start:end], c, start, cache)
+            report["checks"]["cache_max_abs"] = max(errors)
+        del cache, native, actual, full, changed, vc, future
+        model.requires_grad_(True).train()
+        with autocast(device, cfg["precision"]):
+            loss, _ = ARObjective(model, shift=cfg["shift"])(x0, c, chunk_index=0, sigma=0.5)
+        loss.backward()
+        for name in ("dit", "vace", "adapter"):
+            grads = [p.grad for p in getattr(model, name).parameters() if p.grad is not None]
+            if not grads or not all(bool(torch.isfinite(g).all()) for g in grads):
+                raise AssertionError(f"Missing/nonfinite {name} gradients")
+            report["checks"][name + "_grad_abs_sum"] = sum(float(g.abs().sum()) for g in grads)
+        report["checks"]["ar_loss"] = float(loss.detach())
+        report["status"] = "passed"
+    except Exception as error:
+        report["status"] = "failed"
+        report["error"] = repr(error)
+        raise
+    finally:
+        report["elapsed_seconds"] = time.monotonic() - started
+        report["peak_allocated_gib"] = torch.cuda.max_memory_allocated() / 2**30
+        Path(a.output).parent.mkdir(parents=True, exist_ok=True)
+        atomic_json(a.output, report)
+
+
+if __name__ == "__main__":
+    main()

viddub/foundation/NOTICE

@@ -0,0 +1,12 @@
+Reduced DiffSynth-Studio Wan2.1-VACE foundation
+
+Derived from DiffSynth-Studio commit
+7686e54d41d25c0e8ed5f1318acc23b6bb832654 under the Apache License 2.0.
+
+Retained components: Wan DiT, VACE, Wan VAE, Wan T5/tokenizer, gradient
+checkpointing, and the Wan flow-matching schedule. Unrelated model families and
+optional Wan methods were removed. VidDub adds frame-local audio
+cross-attention integration and causal streaming VAE helpers.
+
+The Apache-2.0 license text is available at
+../dependencies/licenses/Apache-2.0.txt.

viddub/foundation/__init__.py

@@ -0,0 +1,8 @@
+"""Minimal Wan2.1-VACE backend used by VidDub.
+
+The retained Wan model code is derived from DiffSynth-Studio commit
+7686e54d41d25c0e8ed5f1318acc23b6bb832654 under Apache-2.0. Unrelated model
+families and Wan extensions are intentionally excluded from this fork.
+"""
+
+from .pipelines.wan_video import ModelConfig, WanVideoPipeline

viddub/foundation/core/__init__.py

@@ -0,0 +1 @@
+from .gradient import gradient_checkpoint_forward

viddub/foundation/core/gradient/__init__.py

@@ -0,0 +1 @@
+from .gradient_checkpoint import gradient_checkpoint_forward