diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..7e6672ab433e43565e1602be2ac0b2408028d6be 100644 --- a/.gitattributes +++ b/.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 diff --git a/README.md b/README.md index 6e8b24f2d9e638066a2b1c379db74f345dfc29d4..5015a092d2e74041768406c302ff1ee4bd915222 100644 --- a/README.md +++ b/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. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..157d076bad5ef4b64731b3f9bce94d629e0765ba --- /dev/null +++ b/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) diff --git a/assets/empty_context.pt b/assets/empty_context.pt new file mode 100644 index 0000000000000000000000000000000000000000..8bb686d6423581d988bee6fd6540ab14f8c3d845 --- /dev/null +++ b/assets/empty_context.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ad9140ae8591e91b64c3d5dbeb183c6701f19d419c8b5f43fe364146042e88ae +size 18003 diff --git a/data/__init__.py b/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1a5486a2d2cea1a71e2afc4e0c5a0331117e2729 --- /dev/null +++ b/data/__init__.py @@ -0,0 +1 @@ +"""VidDub data preparation and preprocessing tools.""" diff --git a/data/common.py b/data/common.py new file mode 100644 index 0000000000000000000000000000000000000000..145244557357df26a640a0dffa519e2d57c15ae7 --- /dev/null +++ b/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] diff --git a/data/s01_analyze.py b/data/s01_analyze.py new file mode 100644 index 0000000000000000000000000000000000000000..a74bdd1a3877ac90da91533c9b2a5cc720c6d97f --- /dev/null +++ b/data/s01_analyze.py @@ -0,0 +1,182 @@ +"""Per video: face landmarks (mediapipe), 16k mono wav, SyncNet offset/confidence. + +Outputs: + data/face//.npz lm [T,468,2] float16 (nan = no face), nfaces [T] + data/audio//.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() diff --git a/data/s02_make_clips.py b/data/s02_make_clips.py new file mode 100644 index 0000000000000000000000000000000000000000..8c1a20201f4422e7cabd11c47f65c1fe281787dc --- /dev/null +++ b/data/s02_make_clips.py @@ -0,0 +1,198 @@ +"""Filter s01 results and cut training clips. + +Per clip (data/clips///): + 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() diff --git a/examples/man_laughing_audio.wav b/examples/man_laughing_audio.wav new file mode 100644 index 0000000000000000000000000000000000000000..a8b694686fec1676f5b8f5a5243484619294a0ec --- /dev/null +++ b/examples/man_laughing_audio.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f53d0999a39251346563a4de62384920c28a67c053ae91fe05c948bf04644c5e +size 144078 diff --git a/examples/man_laughing_video.mp4 b/examples/man_laughing_video.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..d95af7731a53d1de11f8ea7b352d13e721aa12a8 --- /dev/null +++ b/examples/man_laughing_video.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a09134a1e470c253e9241521887c7bd01714eaf45a8a16159b6bd325ec8c0eaa +size 1003686 diff --git a/examples/man_smiling_studio_audio.wav b/examples/man_smiling_studio_audio.wav new file mode 100644 index 0000000000000000000000000000000000000000..c7645f4fa88599e28a7a9669ccc233b8aba30e22 --- /dev/null +++ b/examples/man_smiling_studio_audio.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:da0044fcc67fd95ba6b3e75c8f5507905c046100b8074fb13cb3da82ac6f6c05 +size 144078 diff --git a/examples/man_smiling_studio_video.mp4 b/examples/man_smiling_studio_video.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..1f56bb3cc3c93cf0a6ee14317b3b6fbf7c6a6916 --- /dev/null +++ b/examples/man_smiling_studio_video.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5a54a7fe60ceb0be72dc6cd2f6dd3b91ed3b83db919a0e618cd91168e011f9cb +size 776166 diff --git a/examples/woman_laughing_studio_audio.wav b/examples/woman_laughing_studio_audio.wav new file mode 100644 index 0000000000000000000000000000000000000000..8f683f256d35da76df08ec43894f80e84f2c587f --- /dev/null +++ b/examples/woman_laughing_studio_audio.wav @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:99cb56c5ccb274afeed195232dcd65745ed4eab9559a4b0f34670e32defada7b +size 144078 diff --git a/examples/woman_laughing_studio_video.mp4 b/examples/woman_laughing_studio_video.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..fc14d6a349354e0a2f5f5e1cf9cab98f650c0bf9 --- /dev/null +++ b/examples/woman_laughing_studio_video.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:911cf4085ab10fd231136a1cb7c72e4e2e0cc45da5e8a0672a49924646c26b0c +size 716276 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..fcfd3deea38b31cc222371e8c6674e61bad0ca40 --- /dev/null +++ b/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 diff --git a/viddub/__init__.py b/viddub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..098344d852ae63dc87fcc086025c66b87b7ac41e --- /dev/null +++ b/viddub/__init__.py @@ -0,0 +1,3 @@ +"""VidDub model, training, and inference implementation.""" + +__version__ = "1.0.0" diff --git a/viddub/causal/__init__.py b/viddub/causal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3770fc2f0d52cba6e7980e81d52b05b36af7768e --- /dev/null +++ b/viddub/causal/__init__.py @@ -0,0 +1 @@ +"""Checkpoint-compatible, block-causal Wan-Dub training and sampling.""" diff --git a/viddub/causal/art_alignment.py b/viddub/causal/art_alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..422b412cf0cfe7292d2814f005820635bdb6348b --- /dev/null +++ b/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() diff --git a/viddub/causal/cf/NOTICE b/viddub/causal/cf/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..21ed7caea8a6aa3203d87f10c2a5e2dc5fcda743 --- /dev/null +++ b/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. diff --git a/viddub/causal/cf/__init__.py b/viddub/causal/cf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cedd2ef3d02452870fb2247a0978f24f2d440fe3 --- /dev/null +++ b/viddub/causal/cf/__init__.py @@ -0,0 +1 @@ +"""Causal Forcing++ adaptation for the existing Wan-Dub model and conditions.""" diff --git a/viddub/causal/cf/consistency.py b/viddub/causal/cf/consistency.py new file mode 100644 index 0000000000000000000000000000000000000000..73a2043ae7f9c6980bfcd587936fd8507b380eff --- /dev/null +++ b/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 diff --git a/viddub/causal/cf/dmd.py b/viddub/causal/cf/dmd.py new file mode 100644 index 0000000000000000000000000000000000000000..4e6fc5638f31a594336ec684e8b8147a4e0b7d6e --- /dev/null +++ b/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)} diff --git a/viddub/causal/cf/sampling.py b/viddub/causal/cf/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..a9ee95254a0c2d6ce1fe2b1074f5db7412a196ee --- /dev/null +++ b/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) diff --git a/viddub/causal/cf/train.py b/viddub/causal/cf/train.py new file mode 100644 index 0000000000000000000000000000000000000000..b9703298efb0f516481c81ecb5e5caf415ebea62 --- /dev/null +++ b/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(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] diff --git a/viddub/causal/infer.py b/viddub/causal/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..f72bb389f29021a6656e8c184af69b28953464ac --- /dev/null +++ b/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() diff --git a/viddub/causal/make_trajectories.py b/viddub/causal/make_trajectories.py new file mode 100644 index 0000000000000000000000000000000000000000..4db9ab509be5878eb5a38c60d3999e63bc13c332 --- /dev/null +++ b/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() diff --git a/viddub/causal/model.py b/viddub/causal/model.py new file mode 100644 index 0000000000000000000000000000000000000000..2b082a116f816d67b09e768ea4a54a9062b53c94 --- /dev/null +++ b/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"} diff --git a/viddub/causal/pretrain_fake.py b/viddub/causal/pretrain_fake.py new file mode 100644 index 0000000000000000000000000000000000000000..c006e5bb4ad5286c4f262dae4b21b7d09405bf89 --- /dev/null +++ b/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 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 diff --git a/viddub/causal/runtime.py b/viddub/causal/runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..e83e465dd11bf7a7662508623409a6d1e80c3ecb --- /dev/null +++ b/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 diff --git a/viddub/causal/sampling.py b/viddub/causal/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..2f1f0f71c40006841538f4febe9c31e3a36b26f9 --- /dev/null +++ b/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") diff --git a/viddub/causal/train.py b/viddub/causal/train.py new file mode 100644 index 0000000000000000000000000000000000000000..3138cece795cfd88aaaa18510f626eff223fc701 --- /dev/null +++ b/viddub/causal/train.py @@ -0,0 +1,6 @@ +"""Backward-compatible generic causal training entry point.""" +from .training.engine import main + + +if __name__ == "__main__": + main() diff --git a/viddub/causal/train_ar.py b/viddub/causal/train_ar.py new file mode 100644 index 0000000000000000000000000000000000000000..46c0930b60139e9391e949cce02818f627037e9f --- /dev/null +++ b/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") diff --git a/viddub/causal/train_dmd.py b/viddub/causal/train_dmd.py new file mode 100644 index 0000000000000000000000000000000000000000..d1acc4acf1ccce2b66fe9504c98d1429ae7e2f44 --- /dev/null +++ b/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") diff --git a/viddub/causal/train_ode.py b/viddub/causal/train_ode.py new file mode 100644 index 0000000000000000000000000000000000000000..2aaa38a251b83189f46ecfc5a1964fb20f136a96 --- /dev/null +++ b/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") diff --git a/viddub/causal/training/__init__.py b/viddub/causal/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..930fcd3e8484c4034c80a13fd9d13ac39edab6e0 --- /dev/null +++ b/viddub/causal/training/__init__.py @@ -0,0 +1,3 @@ +"""Shared causal training implementation.""" + +from .stages import ARObjective, DMDObjective, FakeObjective, ODEObjective diff --git a/viddub/causal/training/ar.py b/viddub/causal/training/ar.py new file mode 100644 index 0000000000000000000000000000000000000000..10f8da1eeb95f17bd391f10c3638531d54273605 --- /dev/null +++ b/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 diff --git a/viddub/causal/training/common.py b/viddub/causal/training/common.py new file mode 100644 index 0000000000000000000000000000000000000000..76e126f5d96e7efb0d847b5e4c75872dfc6a82c9 --- /dev/null +++ b/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() diff --git a/viddub/causal/training/dmd.py b/viddub/causal/training/dmd.py new file mode 100644 index 0000000000000000000000000000000000000000..15191ca913a27047de11dfcbd3359850fde641bf --- /dev/null +++ b/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 diff --git a/viddub/causal/training/engine.py b/viddub/causal/training/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..6806ce56a7b6809c4446580c59df4578d74ef2d7 --- /dev/null +++ b/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() diff --git a/viddub/causal/training/fake.py b/viddub/causal/training/fake.py new file mode 100644 index 0000000000000000000000000000000000000000..ad18a94404a54b05e3869193d5df101cb48e9db9 --- /dev/null +++ b/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()} diff --git a/viddub/causal/training/ode.py b/viddub/causal/training/ode.py new file mode 100644 index 0000000000000000000000000000000000000000..04f577adcf38f04872411b38e4795ba525f6a468 --- /dev/null +++ b/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 diff --git a/viddub/causal/training/stages.py b/viddub/causal/training/stages.py new file mode 100644 index 0000000000000000000000000000000000000000..2aa3e592e063eaa999d470c1acd674973e13af44 --- /dev/null +++ b/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 diff --git a/viddub/causal/trajectories.py b/viddub/causal/trajectories.py new file mode 100644 index 0000000000000000000000000000000000000000..d8f5ba65125a77eb0118f2f97eccae699cca52c0 --- /dev/null +++ b/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"]} diff --git a/viddub/causal/verify_checkpoint.py b/viddub/causal/verify_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..1ee5581c4e8a6fe3d94c44db8c98ceb35d9d0399 --- /dev/null +++ b/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() diff --git a/viddub/foundation/NOTICE b/viddub/foundation/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..e56e94f0d51d1ab448448a088f356de5759d7f5c --- /dev/null +++ b/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. diff --git a/viddub/foundation/__init__.py b/viddub/foundation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..32ac339a5d47f1bd93859f95ffd742eaabf82e0f --- /dev/null +++ b/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 diff --git a/viddub/foundation/core/__init__.py b/viddub/foundation/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fb6e76e17a8822714965d877c4fdb2a4219b1d5c --- /dev/null +++ b/viddub/foundation/core/__init__.py @@ -0,0 +1 @@ +from .gradient import gradient_checkpoint_forward diff --git a/viddub/foundation/core/gradient/__init__.py b/viddub/foundation/core/gradient/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..57914792a78ec32f69c3c99ae37535598efc8d52 --- /dev/null +++ b/viddub/foundation/core/gradient/__init__.py @@ -0,0 +1 @@ +from .gradient_checkpoint import gradient_checkpoint_forward diff --git a/viddub/foundation/core/gradient/gradient_checkpoint.py b/viddub/foundation/core/gradient/gradient_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..495827f809d3deab94775de80e2e574fbc3dfaa1 --- /dev/null +++ b/viddub/foundation/core/gradient/gradient_checkpoint.py @@ -0,0 +1,65 @@ +import torch + + +try: + import deepspeed + _HAS_DEEPSPEED = True +except ModuleNotFoundError: + _HAS_DEEPSPEED = False + + +def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + +def create_custom_forward_use_reentrant(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + +def judge_args_requires_grad(*args): + for arg in args: + if isinstance(arg, torch.Tensor) and arg.requires_grad: + return True + return False + + +def gradient_checkpoint_forward( + model, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + *args, + **kwargs, +): + if use_gradient_checkpointing and _HAS_DEEPSPEED and deepspeed.checkpointing.is_configured(): + all_args = args + tuple(kwargs.values()) + if not judge_args_requires_grad(*all_args): + # get the first grad_enabled tensor from un_checkpointed forward + model_output = model(*args, **kwargs) + else: + model_output = deepspeed.checkpointing.checkpoint( + create_custom_forward_use_reentrant(model), + *all_args, + ) + return model_output + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + model_output = torch.utils.checkpoint.checkpoint( + create_custom_forward(model), + *args, + **kwargs, + use_reentrant=False, + ) + elif use_gradient_checkpointing: + model_output = torch.utils.checkpoint.checkpoint( + create_custom_forward(model), + *args, + **kwargs, + use_reentrant=False, + ) + else: + model_output = model(*args, **kwargs) + return model_output diff --git a/viddub/foundation/diffusion/__init__.py b/viddub/foundation/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..744c8dd4bd216c137583dd6f5091d2da731fadf8 --- /dev/null +++ b/viddub/foundation/diffusion/__init__.py @@ -0,0 +1 @@ +from .flow_match import FlowMatchScheduler diff --git a/viddub/foundation/diffusion/flow_match.py b/viddub/foundation/diffusion/flow_match.py new file mode 100644 index 0000000000000000000000000000000000000000..26ae85e86a642dc9195e9dd0105f2aae454bbefb --- /dev/null +++ b/viddub/foundation/diffusion/flow_match.py @@ -0,0 +1,47 @@ +"""Wan flow-matching scheduler retained from DiffSynth-Studio.""" +import torch + + +class FlowMatchScheduler: + def __init__(self, template="Wan"): + if template != "Wan": + raise ValueError("The VidDub fork only supports the Wan schedule") + self.num_train_timesteps = 1000 + self.set_timesteps(100) + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, + training=False, shift=5.0): + sigma_start = float(denoising_strength) + self.sigmas = torch.linspace(sigma_start, 0.0, num_inference_steps + 1)[:-1] + self.sigmas = shift * self.sigmas / (1 + (shift - 1) * self.sigmas) + self.timesteps = self.sigmas * self.num_train_timesteps + self.training = bool(training) + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - 500) / 1000) ** 2) + shifted = y - y.min() + self.linear_timesteps_weights = shifted * (1000 / shifted.sum()) + if len(self.timesteps) != 1000: + self.linear_timesteps_weights *= len(self.timesteps) / 1000 + self.linear_timesteps_weights += self.linear_timesteps_weights[1] + + def _index(self, timestep): + value = timestep.detach().float().cpu() if torch.is_tensor(timestep) else torch.tensor(timestep) + return torch.argmin((self.timesteps - value).abs()) + + def step(self, model_output, timestep, sample, to_final=False): + index = self._index(timestep) + sigma = self.sigmas[index].to(sample) + next_sigma = sample.new_zeros(()) if to_final or index + 1 >= len(self.sigmas) else self.sigmas[index + 1].to(sample) + return sample + model_output * (next_sigma - sigma) + + def add_noise(self, original_samples, noise, timestep): + sigma = self.sigmas[self._index(timestep)].to(original_samples) + return (1 - sigma) * original_samples + sigma * noise + + @staticmethod + def training_target(sample, noise, timestep): + return noise - sample + + def training_weight(self, timestep): + return self.linear_timesteps_weights[self._index(timestep)].to(timestep.device) diff --git a/viddub/foundation/models/__init__.py b/viddub/foundation/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2b75e6a342669eeb2c37d6750e86b4ee2382c616 --- /dev/null +++ b/viddub/foundation/models/__init__.py @@ -0,0 +1,4 @@ +from .wan_video_dit import WanModel +from .wan_video_text_encoder import HuggingfaceTokenizer, WanTextEncoder +from .wan_video_vace import VaceWanModel +from .wan_video_vae import WanVideoVAE diff --git a/viddub/foundation/models/wan_video_dit.py b/viddub/foundation/models/wan_video_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..0aeba656d665bc8f89c2398b8059abc1c2042e00 --- /dev/null +++ b/viddub/foundation/models/wan_video_dit.py @@ -0,0 +1,406 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +from typing import Tuple, Optional +from einops import rearrange +from ..core.gradient import gradient_checkpoint_forward + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + from sageattention import sageattn + SAGE_ATTN_AVAILABLE = True +except ModuleNotFoundError: + SAGE_ATTN_AVAILABLE = False + + +def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False): + if torch.is_autocast_enabled(): + # fp32 weights (full finetune) make RMSNorm return fp32 q/k, which flash-attn rejects. + dtype = torch.get_autocast_dtype("cuda") + q, k, v = q.to(dtype), k.to(dtype), v.to(dtype) + if compatibility_mode: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + elif q.is_cuda and FLASH_ATTN_3_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x,tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif q.is_cuda and FLASH_ATTN_2_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn.flash_attn_func(q, k, v) + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif q.is_cuda and SAGE_ATTN_AVAILABLE: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = sageattn(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + else: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + return x + + +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return (x * (1 + scale) + shift) + + +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer(position.type(torch.float64), torch.pow( + 10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + + +def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) + return f_freqs_cis, h_freqs_cis, w_freqs_cis + + +def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2) + [: (dim // 2)].double() / dim)) + freqs = torch.outer(torch.arange(end, device=freqs.device), freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + +def rope_apply(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float64).reshape( + x.shape[0], x.shape[1], x.shape[2], -1, 2)) + freqs = freqs.to(torch.complex64) if freqs.device.type == "npu" else freqs + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +def set_to_torch_norm(models): + for model in models: + for module in model.modules(): + if isinstance(module, RMSNorm): + module.use_torch_norm = True + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + self.use_torch_norm = False + self.normalized_shape = (dim,) + + def norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def forward(self, x): + dtype = x.dtype + if self.use_torch_norm: + return F.rms_norm(x, self.normalized_shape, self.weight, self.eps) + else: + return self.norm(x.float()).to(dtype) * self.weight + + +class AttentionModule(nn.Module): + def __init__(self, num_heads): + super().__init__() + self.num_heads = num_heads + + def forward(self, q, k, v): + x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads) + return x + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x, freqs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + x = self.attn(q, k, v) + return self.o(x) + + +class CrossAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + self.has_image_input = has_image_input + if has_image_input: + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + self.norm_k_img = RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + if self.has_image_input: + img = y[:, :257] + ctx = y[:, 257:] + else: + ctx = y + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(ctx)) + v = self.v(ctx) + x = self.attn(q, k, v) + if self.has_image_input: + k_img = self.norm_k_img(self.k_img(img)) + v_img = self.v_img(img) + y = flash_attention(q, k_img, v_img, num_heads=self.num_heads) + x = x + y + return self.o(x) + + +class GateModule(nn.Module): + def __init__(self,): + super().__init__() + + def forward(self, x, gate, residual): + return x + gate * residual + +class DiTBlock(nn.Module): + def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention( + dim, num_heads, eps, has_image_input=has_image_input) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU( + approximate='tanh'), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + + def forward(self, x, context, t_mod, freqs): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + # msa: multi-head self-attention mlp: multi-layer perceptron + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim) + if has_seq: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), + shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), + ) + input_x = modulate(self.norm1(x), shift_msa, scale_msa) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x), context) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class MLP(torch.nn.Module): + def __init__(self, in_dim, out_dim, has_pos_emb=False): + super().__init__() + self.proj = torch.nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, in_dim), + nn.GELU(), + nn.Linear(in_dim, out_dim), + nn.LayerNorm(out_dim) + ) + self.has_pos_emb = has_pos_emb + if has_pos_emb: + self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280))) + + def forward(self, x): + if self.has_pos_emb: + x = x + self.emb_pos.to(dtype=x.dtype, device=x.device) + return self.proj(x) + + +class Head(nn.Module): + def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + if len(t_mod.shape) == 3: + shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2) + x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2))) + else: + shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + scale) + shift)) + return x + + +class WanModel(torch.nn.Module): + + _repeated_blocks = ["DiTBlock"] + + def __init__( + self, + dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + patch_size: Tuple[int, int, int], + num_heads: int, + num_layers: int, + has_image_input: bool, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + seperated_timestep: bool = False, + require_vae_embedding: bool = True, + require_clip_embedding: bool = True, + fuse_vae_embedding_in_latents: bool = False, + ): + super().__init__() + self.dim = dim + self.in_dim = in_dim + self.freq_dim = freq_dim + self.has_image_input = has_image_input + self.patch_size = patch_size + self.seperated_timestep = seperated_timestep + self.require_vae_embedding = require_vae_embedding + self.require_clip_embedding = require_clip_embedding + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate='tanh'), + nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList([ + DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps) + for _ in range(num_layers) + ]) + self.head = Head(dim, out_dim, patch_size, eps) + head_dim = dim // num_heads + + self.freqs = precompute_freqs_cis_3d(head_dim) + + if has_image_input: + self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb) # clip_feature_dim = 1280 + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + def patchify(self, x: torch.Tensor): + return self.patch_embedding(x) + + def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): + return rearrange( + x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', + f=grid_size[0], h=grid_size[1], w=grid_size[2], + x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2] + ) + + def forward(self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ): + t = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep).to(x.dtype)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + if self.has_image_input: + x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) + clip_embdding = self.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + + x = self.patchify(x) + f, h, w = x.shape[2:] + x = rearrange(x, "b c f h w -> b (f h w) c").contiguous() + + freqs = torch.cat([ + self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) + + for block in self.blocks: + if self.training: + x = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + x, context, t_mod, freqs + ) + else: + x = block(x, context, t_mod, freqs) + + x = self.head(x, t) + x = self.unpatchify(x, (f, h, w)) + return x diff --git a/viddub/foundation/models/wan_video_text_encoder.py b/viddub/foundation/models/wan_video_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..64090db8c65138abfdb60a822b3ba2e74fefeb4c --- /dev/null +++ b/viddub/foundation/models/wan_video_text_encoder.py @@ -0,0 +1,330 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import AutoTokenizer +import ftfy +import html +import string +import regex as re + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + + def __init__(self, dim, eps=1e-6): + super(T5LayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super(T5Attention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + # layers + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C]. + context: [B, L2, C] or None. + mask: [B, L2] or [B, L1, L2] or None. + """ + # check inputs + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + # attention bias + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in [2, 3] + mask = mask.view(b, 1, 1, + -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # compute attention (T5 does not use scaling) + attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + + # output + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + + def __init__(self, dim, dim_ffn, dropout=0.1): + super(T5FeedForward, self).__init__() + self.dim = dim + self.dim_ffn = dim_ffn + + # layers + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5SelfAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5SelfAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + return x + + +class T5RelativeEmbedding(nn.Module): + + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super(T5RelativeEmbedding, self).__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + + # layers + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \ + # torch.arange(lq).unsqueeze(1).to(device) + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ + torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( + 0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + # preprocess + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + # embeddings for small and large positions + max_exact = num_buckets // 2 + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / + math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() + rel_pos_large = torch.min( + rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_( + m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5) + + +class WanTextEncoder(torch.nn.Module): + + def __init__(self, + vocab=256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1): + super(WanTextEncoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \ + else nn.Embedding(vocab, dim) + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) if shared_pos else None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + e = self.pos_embedding(x.size(1), + x.size(1)) if self.shared_pos else None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +def canonicalize(text, keep_punctuation_exact_string=None): + text = text.replace('_', ' ') + if keep_punctuation_exact_string: + text = keep_punctuation_exact_string.join( + part.translate(str.maketrans('', '', string.punctuation)) + for part in text.split(keep_punctuation_exact_string)) + else: + text = text.translate(str.maketrans('', '', string.punctuation)) + text = text.lower() + text = re.sub(r'\s+', ' ', text) + return text.strip() + + +class HuggingfaceTokenizer: + + def __init__(self, name, seq_len=None, clean=None, **kwargs): + assert clean in (None, 'whitespace', 'lower', 'canonicalize') + self.name = name + self.seq_len = seq_len + self.clean = clean + + # init tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs) + self.vocab_size = self.tokenizer.vocab_size + + def __call__(self, sequence, **kwargs): + return_mask = kwargs.pop('return_mask', False) + + # arguments + _kwargs = {'return_tensors': 'pt'} + if self.seq_len is not None: + _kwargs.update({ + 'padding': 'max_length', + 'truncation': True, + 'max_length': self.seq_len + }) + _kwargs.update(**kwargs) + + # tokenization + if isinstance(sequence, str): + sequence = [sequence] + if self.clean: + sequence = [self._clean(u) for u in sequence] + ids = self.tokenizer(sequence, **_kwargs) + + # output + if return_mask: + return ids.input_ids, ids.attention_mask + else: + return ids.input_ids + + def _clean(self, text): + if self.clean == 'whitespace': + text = whitespace_clean(basic_clean(text)) + elif self.clean == 'lower': + text = whitespace_clean(basic_clean(text)).lower() + elif self.clean == 'canonicalize': + text = canonicalize(basic_clean(text)) + return text \ No newline at end of file diff --git a/viddub/foundation/models/wan_video_vace.py b/viddub/foundation/models/wan_video_vace.py new file mode 100644 index 0000000000000000000000000000000000000000..0e13183c233f58974b920e5582aa7631217abb4a --- /dev/null +++ b/viddub/foundation/models/wan_video_vace.py @@ -0,0 +1,74 @@ +import torch +from .wan_video_dit import DiTBlock +from ..core.gradient import gradient_checkpoint_forward + +class VaceWanAttentionBlock(DiTBlock): + def __init__(self, has_image_input, dim, num_heads, ffn_dim, eps=1e-6, block_id=0): + super().__init__(has_image_input, dim, num_heads, ffn_dim, eps=eps) + self.block_id = block_id + if block_id == 0: + self.before_proj = torch.nn.Linear(self.dim, self.dim) + self.after_proj = torch.nn.Linear(self.dim, self.dim) + + def forward(self, c, x, context, t_mod, freqs): + if self.block_id == 0: + c = self.before_proj(c) + x + all_c = [] + else: + all_c = list(torch.unbind(c)) + c = all_c.pop(-1) + c = super().forward(c, context, t_mod, freqs) + c_skip = self.after_proj(c) + all_c += [c_skip, c] + c = torch.stack(all_c) + return c + + +class VaceWanModel(torch.nn.Module): + def __init__( + self, + vace_layers=(0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28), + vace_in_dim=96, + patch_size=(1, 2, 2), + has_image_input=False, + dim=1536, + num_heads=12, + ffn_dim=8960, + eps=1e-6, + ): + super().__init__() + self.vace_layers = vace_layers + self.vace_in_dim = vace_in_dim + self.vace_layers_mapping = {i: n for n, i in enumerate(self.vace_layers)} + + # vace blocks + self.vace_blocks = torch.nn.ModuleList([ + VaceWanAttentionBlock(has_image_input, dim, num_heads, ffn_dim, eps, block_id=i) + for i in self.vace_layers + ]) + + # vace patch embeddings + self.vace_patch_embedding = torch.nn.Conv3d(vace_in_dim, dim, kernel_size=patch_size, stride=patch_size) + + def forward( + self, x, vace_context, context, t_mod, freqs, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + ): + c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context] + c = [u.flatten(2).transpose(1, 2) for u in c] + c = torch.cat([ + torch.cat([u, u.new_zeros(1, x.shape[1] - u.size(1), u.size(2))], + dim=1) for u in c + ]) + + for block in self.vace_blocks: + c = gradient_checkpoint_forward( + block, + use_gradient_checkpointing, + use_gradient_checkpointing_offload, + c, x, context, t_mod, freqs + ) + + hints = torch.unbind(c)[:-1] + return hints diff --git a/viddub/foundation/models/wan_video_vae.py b/viddub/foundation/models/wan_video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..4fdd5df7e25ddd1e04af320ba3199fd058950eff --- /dev/null +++ b/viddub/foundation/models/wan_video_vae.py @@ -0,0 +1,1456 @@ +from einops import rearrange, repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +CACHE_T = 2 + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +def block_causal_mask(x, block_size): + # params + b, n, s, _, device = *x.size(), x.device + assert s % block_size == 0 + num_blocks = s // block_size + + # build mask + mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=device) + for i in range(num_blocks): + mask[:, :, + i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 + return mask + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x, feat_cache, feat_idx + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size) + return x + + +class Resample38(Resample): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super(Resample, self).__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + else: + self.resample = nn.Identity() + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h, feat_cache, feat_idx + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + #attn_mask=block_causal_mask(q, block_size=h * w) + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample38(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x, feat_cache, feat_idx = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy), feat_cache, feat_idx + + +class Up_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False + ): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample38(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main, feat_cache, feat_idx = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut, feat_cache, feat_idx + else: + return x_main, feat_cache, feat_idx + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x, feat_cache, feat_idx + + +class Encoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] if i < len(temperal_downsample) else False + ) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x, feat_cache, feat_idx + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x, feat_cache, feat_idx + + + +class Decoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock(in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1)) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1)) + + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x, feat_cache, feat_idx = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :] + .unsqueeze(2) + .to(cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x, feat_cache, feat_idx + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class VideoVAE_(nn.Module): + + def __init__(self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + def encode(self, x, scale): + self.clear_cache() + ## cache + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + self.clear_cache() + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + else: + out_, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) # may add tensor offload + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +class WanVideoVAE(nn.Module): + + def __init__(self, z_dim=16): + super().__init__() + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) + self.upsampling_factor = 8 + self.z_dim = z_dim + + + def build_1d_mask(self, length, left_bound, right_bound, border_width): + x = torch.ones((length,)) + if not left_bound: + x[:border_width] = (torch.arange(border_width) + 1) / border_width + if not right_bound: + x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) + return x + + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) + + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + + mask = torch.stack([h, w]).min(dim=0).values + mask = rearrange(mask, "H W -> 1 1 1 H W") + return mask + + + def tiled_decode(self, hidden_states, device, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + data_device = "cpu" + computation_device = device + + out_T = T * 4 - 3 + weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) + values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): + hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device) + hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype, device=data_device) + + target_h = h * self.upsampling_factor + target_w = w * self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.clamp_(-1, 1) + return values + + + def tiled_encode(self, video, device, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + data_device = "cpu" + computation_device = device + + out_T = (T + 3) // 4 + weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): + hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device) + hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype, device=data_device) + + target_h = h // self.upsampling_factor + target_w = w // self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + return values + + + def single_encode(self, video, device): + video = video.to(device) + x = self.model.encode(video, self.scale) + return x + + + def streaming_encode(self, frame_chunks, device): + """Causally encode one 1-frame chunk followed by 4-frame chunks. + + This is algebraically the same temporal loop used by ``VideoVAE_.encode`` + but avoids moving a full long video to the accelerator. Spatial shape, + dtype and batch must stay fixed. Returned latent chunks each contain one + temporal latent and may be consumed immediately or offloaded to CPU. + """ + self.model.clear_cache() + expected = 1 + reference = None + for index, chunk in enumerate(frame_chunks): + if chunk.ndim != 5 or chunk.shape[2] != expected: + raise ValueError(f"Streaming VAE chunk {index} must contain {expected} frames") + signature = (chunk.shape[0], chunk.shape[1], chunk.shape[3], chunk.shape[4], chunk.dtype) + if reference is None: reference = signature + elif signature != reference: raise ValueError("Streaming VAE geometry changed") + self.model._enc_conv_idx = [0] + out, self.model._enc_feat_map, self.model._enc_conv_idx = self.model.encoder( + chunk.to(device), feat_cache=self.model._enc_feat_map, + feat_idx=self.model._enc_conv_idx) + mu, _ = self.model.conv1(out).chunk(2, dim=1) + scale = self.scale + if isinstance(scale[0], torch.Tensor): + values = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu-values[0].view(1,self.z_dim,1,1,1))*values[1].view(1,self.z_dim,1,1,1) + else: + values = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu-values[0])*values[1] + yield mu + expected = 4 + + + def streaming_decode(self, latent_chunks, device): + """Causally decode one-latent chunks without storing all latents on GPU.""" + self.model.clear_cache() + reference = None + for index, latent in enumerate(latent_chunks): + if latent.ndim != 5 or latent.shape[2] != 1: + raise ValueError(f"Streaming decode chunk {index} must contain one latent frame") + signature = (latent.shape[0], latent.shape[1], latent.shape[3], latent.shape[4], latent.dtype) + if reference is None: reference = signature + elif signature != reference: raise ValueError("Streaming decode geometry changed") + scale = self.scale + z = latent.to(device) + if isinstance(scale[0], torch.Tensor): + values = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z/values[1].view(1,self.z_dim,1,1,1)+values[0].view(1,self.z_dim,1,1,1) + else: + values = scale.to(dtype=z.dtype, device=z.device) + z = z/values[1]+values[0] + x = self.model.conv2(z) + self.model._conv_idx = [0] + out, self.model._feat_map, self.model._conv_idx = self.model.decoder( + x, feat_cache=self.model._feat_map, feat_idx=self.model._conv_idx) + yield out.clamp_(-1,1) + + + def single_decode(self, hidden_state, device): + hidden_state = hidden_state.to(device) + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + + def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + videos = [video.to("cpu") for video in videos] + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, device, tile_size, tile_stride) + else: + hidden_state = self.single_encode(video, device) + hidden_state = hidden_state.squeeze(0) + hidden_states.append(hidden_state) + hidden_states = torch.stack(hidden_states) + return hidden_states + + + def decode(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states] + videos = [] + for hidden_state in hidden_states: + hidden_state = hidden_state.unsqueeze(0) + if tiled: + video = self.tiled_decode(hidden_state, device, tile_size, tile_stride) + else: + video = self.single_decode(hidden_state, device) + video = video.squeeze(0) + videos.append(video) + videos = torch.stack(videos) + return videos + + + def encode_framewise(self, videos, device): + hidden_states = [] + for i in range(videos.shape[2]): + hidden_states.append(self.single_encode(videos[:, :, i:i+1], device)) + hidden_states = torch.concat(hidden_states, dim=2) + return hidden_states + + + def decode_framewise(self, hidden_states, device): + video = [] + for i in range(hidden_states.shape[2]): + video.append(self.single_decode(hidden_states[:, :, i:i+1], device)) + video = torch.concat(video, dim=2) + return video + + + @staticmethod + def state_dict_converter(): + return WanVideoVAEStateDictConverter() + + +class WanVideoVAEStateDictConverter: + + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_ + + +class VideoVAE38_(VideoVAE_): + + def __init__(self, + dim=160, + z_dim=48, + dec_dim=256, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super(VideoVAE_, self).__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True) + else: + out_, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + +class WanVideoVAE38(WanVideoVAE): + + def __init__(self, z_dim=48, dim=160): + super(WanVideoVAE, self).__init__() + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE38_(z_dim=z_dim, dim=dim).eval().requires_grad_(False) + self.upsampling_factor = 16 + self.z_dim = z_dim diff --git a/viddub/foundation/pipelines/__init__.py b/viddub/foundation/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..280efdd7e7bdd52bc7ba6b7c3fc338a787aed86a --- /dev/null +++ b/viddub/foundation/pipelines/__init__.py @@ -0,0 +1 @@ +from .wan_video import ModelConfig, WanVideoPipeline, WanVideoUnit_PromptEmbedder, model_fn_wan_video diff --git a/viddub/foundation/pipelines/wan_video.py b/viddub/foundation/pipelines/wan_video.py new file mode 100644 index 0000000000000000000000000000000000000000..1ed17803f14bd13a51e6ff8d0d3e4349e68ee7c7 --- /dev/null +++ b/viddub/foundation/pipelines/wan_video.py @@ -0,0 +1,182 @@ +"""Minimal Wan2.1-VACE-1.3B pipeline for VidDub. + +This module intentionally exposes only the interfaces used by VidDub. It is +derived from DiffSynth-Studio commit 7686e54d41d25c0e8ed5f1318acc23b6bb832654 +under Apache-2.0; unrelated model families and Wan extensions are removed. +""" +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import torch +from einops import rearrange +from safetensors import safe_open + +from ..core import gradient_checkpoint_forward +from ..diffusion import FlowMatchScheduler +from ..models.wan_video_dit import WanModel, precompute_freqs_cis_3d, sinusoidal_embedding_1d +from ..models.wan_video_text_encoder import HuggingfaceTokenizer, WanTextEncoder +from ..models.wan_video_vace import VaceWanModel +from ..models.wan_video_vae import WanVideoVAE + + +@dataclass +class ModelConfig: + path: str + + +def _load_torch(path): + state = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + if len(state) == 1 and next(iter(state)) in ("state_dict", "model_state", "module"): + state = state[next(iter(state))] + return state + + +def _load_safetensors_subset(path, predicate): + state = {} + with safe_open(str(path), framework="pt", device="cpu") as handle: + for key in handle.keys(): + if predicate(key): + state[key] = handle.get_tensor(key) + return state + + +def _load_assign(model, state): + """Materialize a meta-initialized module directly from one state dict.""" + missing, unexpected = model.load_state_dict(state, strict=False, assign=True) + if missing or unexpected: + raise RuntimeError( + f"Checkpoint mismatch for {type(model).__name__}: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}") + return model + + +class WanVideoPipeline(torch.nn.Module): + def __init__(self, device="cpu", torch_dtype=torch.bfloat16): + super().__init__() + self.device = torch.device(device) + self.torch_dtype = torch_dtype + self.scheduler = FlowMatchScheduler("Wan") + self.dit = self.vace = self.vae = self.text_encoder = None + self.tokenizer = None + + @classmethod + def from_pretrained(cls, torch_dtype=torch.bfloat16, device="cpu", + model_configs=None, tokenizer_config=None, **_): + configs = [Path(config.path) for config in (model_configs or [])] + diffusion = next((p for p in configs if p.suffix == ".safetensors"), None) + vae_path = next((p for p in configs if "VAE" in p.name.upper()), None) + text_path = next((p for p in configs if "T5" in p.name.upper()), None) + if diffusion is None or vae_path is None: + raise ValueError("Wan2.1-VACE diffusion and VAE checkpoints are required") + + pipe = cls(device=device, torch_dtype=torch_dtype) + with torch.device("meta"): + pipe.dit = WanModel( + has_image_input=False, patch_size=(1, 2, 2), in_dim=16, + dim=1536, ffn_dim=8960, freq_dim=256, text_dim=4096, + out_dim=16, num_heads=12, num_layers=30, eps=1e-6) + pipe.vace = VaceWanModel() + composite = _load_safetensors_subset(diffusion, lambda key: not key.startswith("vace")) + _load_assign(pipe.dit, composite) + pipe.dit.freqs = precompute_freqs_cis_3d(pipe.dit.dim // 12) + del composite + composite = _load_safetensors_subset(diffusion, lambda key: key.startswith("vace")) + _load_assign(pipe.vace, composite) + del composite + + with torch.device("meta"): + pipe.vae = WanVideoVAE() + vae_state = _load_torch(vae_path) + if not any(key.startswith("model.") for key in vae_state): + vae_state = {"model." + key: value for key, value in vae_state.items()} + _load_assign(pipe.vae, vae_state) + mean = torch.tensor([ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921]) + std = torch.tensor([ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160]) + pipe.vae.mean, pipe.vae.std = mean, std + pipe.vae.scale = [mean, 1.0 / std] + del vae_state + + if text_path is not None: + with torch.device("meta"): + pipe.text_encoder = WanTextEncoder() + _load_assign(pipe.text_encoder, _load_torch(text_path)) + if tokenizer_config is not None: + pipe.tokenizer = HuggingfaceTokenizer( + tokenizer_config.path, seq_len=512, clean="whitespace") + pipe.to(device=device, dtype=torch_dtype) + return pipe + + def to(self, *args, **kwargs): + device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs) + if device is not None: + self.device = device + if dtype is not None: + self.torch_dtype = dtype + return super().to(*args, **kwargs) + + def preprocess_image(self, image, min_value=-1, max_value=1): + array = torch.as_tensor(__import__("numpy").array(image), device=self.device) + return array.to(self.torch_dtype).permute(2, 0, 1)[None] * ((max_value-min_value)/255) + min_value + + def preprocess_video(self, video, min_value=-1, max_value=1): + return torch.stack([self.preprocess_image(frame, min_value, max_value)[0] + for frame in video], dim=1)[None] + + +class WanVideoUnit_PromptEmbedder: + def encode_prompt(self, pipe, prompt): + if pipe.text_encoder is None or pipe.tokenizer is None: + raise ValueError("Text encoder/tokenizer is not loaded") + ids, mask = pipe.tokenizer(prompt, return_mask=True, add_special_tokens=True) + ids, mask = ids.to(pipe.device), mask.to(pipe.device) + embedding = pipe.text_encoder(ids, mask) + for index, length in enumerate(mask.gt(0).sum(1).long()): + embedding[index, length:] = 0 + return embedding + + +def model_fn_wan_video( + dit: WanModel, + vace: Optional[VaceWanModel] = None, + latents: torch.Tensor = None, + timestep: torch.Tensor = None, + context: torch.Tensor = None, + vace_context: torch.Tensor = None, + vace_scale: float = 1.0, + audio_adapter=None, + audio_tokens: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **_, +): + t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep)) + t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim)) + context = dit.text_embedding(context) + x = dit.patchify(latents) + f, h, w = x.shape[2:] + x = rearrange(x, "b c f h w -> b (f h w) c").contiguous() + freqs = torch.cat([ + dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], -1).reshape(f * h * w, 1, -1).to(x.device) + hints = None if vace_context is None else vace( + x, vace_context, context, t_mod, freqs, + use_gradient_checkpointing=use_gradient_checkpointing, + use_gradient_checkpointing_offload=use_gradient_checkpointing_offload) + for block_id, block in enumerate(dit.blocks): + x = gradient_checkpoint_forward( + block, use_gradient_checkpointing, + use_gradient_checkpointing_offload, x, context, t_mod, freqs) + if hints is not None and block_id in vace.vace_layers_mapping: + x = x + hints[vace.vace_layers_mapping[block_id]] * vace_scale + if audio_adapter is not None and audio_tokens is not None: + x = gradient_checkpoint_forward( + audio_adapter.blocks[block_id], use_gradient_checkpointing, + use_gradient_checkpointing_offload, x, audio_tokens, f) + return dit.unpatchify(dit.head(x, t), (f, h, w)) diff --git a/viddub/full_context/__init__.py b/viddub/full_context/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ee33a020a2ebda8b668c88e8f3a69b4f1386d5d4 --- /dev/null +++ b/viddub/full_context/__init__.py @@ -0,0 +1 @@ +"""Full-context teacher training and video dubbing utilities.""" diff --git a/viddub/full_context/audio.py b/viddub/full_context/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..d9715a118b2b6a1a3a4ca5322704342cb4d9a385 --- /dev/null +++ b/viddub/full_context/audio.py @@ -0,0 +1,138 @@ +"""Audio conditioning for Wan-Dub: frozen Whisper encoder -> per-latent-frame tokens -> frame-local cross-attn in every DiT block.""" +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import WhisperFeatureExtractor, WhisperModel + +WHISPER_HZ = 50 +VIDEO_FPS = 24 +TOKENS_PER_LATENT = 16 # audio frames (20ms each) around each latent frame's centre = 320ms + + +def stitch_hidden_segments(segments, total_frames): + """Blend overlapping 50-Hz hidden-state segments on one global timeline. + + Each item is ``(global_start_frame, hidden_states)`` with hidden states + shaped [B,L,T,C]. Linear edge ramps suppress segment-boundary context + discontinuities. Exact source values are reconstructed when overlapping + segments contain the same global features. + """ + if total_frames < 1 or not segments: + raise ValueError("Need nonempty segments and a positive output length") + first = segments[0][1] + if first.ndim != 4: + raise ValueError("Expected [B,L,T,C] hidden states") + values = first.new_zeros((*first.shape[:2], total_frames, first.shape[3]), dtype=torch.float32) + weights = first.new_zeros((1, 1, total_frames, 1), dtype=torch.float32) + for start, hidden in segments: + if hidden.shape[:2] != first.shape[:2] or hidden.shape[3] != first.shape[3]: + raise ValueError("Hidden-state segment shapes differ") + if start < 0 or start >= total_frames: + raise ValueError("Segment start is outside the global timeline") + length = min(hidden.shape[2], total_frames-start) + if length < 1: continue + # Positive triangular weights; a single segment therefore remains exact. + position = torch.arange(length, device=hidden.device, dtype=torch.float32) + edge = torch.minimum(position+1, torch.tensor(float(length),device=hidden.device)-position) + weight = edge.view(1,1,length,1) + values[:,:,start:start+length] += hidden[:,:,:length].float()*weight + weights[:,:,start:start+length] += weight + if bool((weights == 0).any()): + raise ValueError("Segment stitching left uncovered audio frames") + return (values/weights).to(first.dtype) + + +class WhisperFeatures(nn.Module): + """Frozen whisper encoder; returns all hidden states [B, L, T_a, C] for the first `seconds` of audio.""" + + def __init__(self, path): + super().__init__() + self.extractor = WhisperFeatureExtractor.from_pretrained(path) + self.encoder = WhisperModel.from_pretrained(path).encoder.eval().requires_grad_(False) + + @torch.no_grad() + def forward(self, audios, device, dtype): + feats = self.extractor(audios, sampling_rate=16000, return_tensors="pt").input_features.to(device, dtype) + hs = self.encoder(feats, output_hidden_states=True).hidden_states # (L+1) x [B, 1500, C] + n = math.ceil(max(len(a) for a in audios) / 16000 * WHISPER_HZ) + return torch.stack([h[:, :n] for h in hs], 1) + + @torch.no_grad() + def forward_long(self, audio, device, dtype, window_seconds=30, hop_seconds=20): + """Encode arbitrarily long mono 16-kHz audio with overlapping windows. + + Audio no longer than one Whisper window uses ``forward`` directly and is + bitwise identical to the existing path. Longer audio is encoded in + overlapping windows and blended on the global 50-Hz timeline. + """ + if getattr(audio, "ndim", 1) != 1: + raise ValueError("forward_long expects one mono waveform") + if not 0 < hop_seconds <= window_seconds <= 30: + raise ValueError("Require 0 < hop <= window <= 30 seconds") + window = int(window_seconds*16000); hop = int(hop_seconds*16000) + if len(audio) <= window: + return self.forward([audio], device, dtype) + total = math.ceil(len(audio)/16000*WHISPER_HZ) + starts = list(range(0, max(1, len(audio)-window+1), hop)) + last = max(0, len(audio)-window) + if not starts or starts[-1] != last: starts.append(last) + segments = [] + for sample_start in starts: + waveform = audio[sample_start:min(len(audio),sample_start+window)] + hidden = self.forward([waveform], device, dtype) + frame_start = round(sample_start/16000*WHISPER_HZ) + segments.append((frame_start, hidden)) + return stitch_hidden_segments(segments, total) + + +def latent_frame_centres(num_frames, fps=VIDEO_FPS): + """Centre time (s) of each VAE latent frame: latent 0 = frame 0, latent k = frames 4k-3..4k.""" + n_lat = (num_frames - 1) // 4 + 1 + return torch.tensor([0.0] + [(4 * k - 1.5) / fps for k in range(1, n_lat)]) + + +class AudioCrossAttn(nn.Module): + def __init__(self, dim, num_heads): + super().__init__() + self.num_heads = num_heads + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.q = nn.Linear(dim, dim) + self.kv = nn.Linear(dim, 2 * dim) + self.o = nn.Linear(dim, dim) + nn.init.zeros_(self.o.weight) + nn.init.zeros_(self.o.bias) + + def forward(self, x, audio_tokens, num_latent_frames): + """x: [B, F*HW, C]; audio_tokens: [B, F, K, C]. Each latent frame only attends to its own audio window.""" + B, N, C = x.shape + Fr, K = audio_tokens.shape[1], audio_tokens.shape[2] + assert Fr == num_latent_frames, (Fr, num_latent_frames) + h = self.num_heads + q = self.q(self.norm(x)).view(B * Fr, N // Fr, h, C // h).transpose(1, 2) + k, v = self.kv(audio_tokens).view(B * Fr, K, 2, h, C // h).permute(2, 0, 3, 1, 4) + out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, C) + return x + self.o(out) + + +class DubAudioAdapter(nn.Module): + def __init__(self, whisper_layers=13, whisper_dim=768, dim=1536, num_heads=12, num_blocks=30): + super().__init__() + self.layer_logits = nn.Parameter(torch.zeros(whisper_layers)) + self.proj = nn.Sequential(nn.Linear(whisper_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim), nn.LayerNorm(dim)) + self.pos = nn.Parameter(torch.zeros(TOKENS_PER_LATENT, dim)) + self.null = nn.Parameter(torch.zeros(TOKENS_PER_LATENT, dim)) + self.blocks = nn.ModuleList([AudioCrossAttn(dim, num_heads) for _ in range(num_blocks)]) + + def tokens(self, whisper_hs, num_frames, num_ref_frames, drop, fps=VIDEO_FPS): + """whisper_hs: [B, L, T_a, Cw] -> [B, num_ref + n_lat, K, dim]. drop: bool tensor [B] -> use null tokens.""" + B, _, T_a, _ = whisper_hs.shape + feats = (whisper_hs * self.layer_logits.softmax(0).view(1, -1, 1, 1).to(whisper_hs.dtype)).sum(1) # B, T_a, Cw + centres = latent_frame_centres(num_frames, fps) * WHISPER_HZ + idx = (centres[:, None] + torch.arange(TOKENS_PER_LATENT) - TOKENS_PER_LATENT / 2).round().long().clamp(0, T_a - 1) + tok = self.proj(feats[:, idx.to(feats.device)]) + self.pos.to(feats.dtype) # B, n_lat, K, dim + null = self.null.to(tok.dtype).expand_as(tok) + tok = torch.where(drop.view(B, 1, 1, 1).to(tok.device), null, tok) + ref = self.null.to(tok.dtype).expand(B, num_ref_frames, -1, -1) + return torch.cat([ref, tok], 1) diff --git a/viddub/full_context/check_equivalence.py b/viddub/full_context/check_equivalence.py new file mode 100644 index 0000000000000000000000000000000000000000..af07df7691787a02e1cd7027269c22ffc368601e --- /dev/null +++ b/viddub/full_context/check_equivalence.py @@ -0,0 +1,31 @@ +"""One-off: cached-latent path (compute_latents + build_inputs) must match DiffSynth's own units.""" +import numpy as np +import torch +from PIL import Image + +from .runtime import ZeroLatents, build_inputs, compute_latents, load_clip, load_models, load_rows + +pipe, _ = load_models("cuda") +pipe.scheduler.set_timesteps(1000, training=True) +clip = load_clip(load_rows("train")[0], 49) + +masked = clip["frames"].copy() +masked[:, clip["mask"]] = 128 +mask_img = Image.fromarray(clip["mask"].astype(np.uint8) * 255).convert("RGB") +shared = {"input_video": [Image.fromarray(f) for f in clip["frames"]], "vace_video": [Image.fromarray(f) for f in masked], + "vace_video_mask": [mask_img] * 49, "vace_reference_image": Image.fromarray(clip["ref"]), "vace_scale": 1, + "height": 512, "width": 512, "num_frames": 49, "seed": 0, "rand_device": "cpu", "cfg_scale": 1, "cfg_merge": False, "tiled": False} +posi, nega = {"prompt": clip["prompt"], "positive": True}, {} +with torch.no_grad(): + for unit in pipe.units: + shared, posi, nega = pipe.unit_runner(unit, pipe, shared, posi, nega) + lat = compute_latents(pipe, clip) + x0, vace, ctx = build_inputs(lat, clip["mask"], 49, ZeroLatents(pipe), context=lat["context"]) + +for name, a, b in (("x0", shared["input_latents"], x0), ("vace_context", shared["vace_context"], vace), ("context", posi["context"], ctx)): + a, b = a.float(), b.float() + print(f"{name}: shapes {tuple(a.shape)} {tuple(b.shape)} max|diff| {(a - b).abs().max().item():.4f} mean|a| {a.abs().mean().item():.4f}") +a, b = shared["vace_context"].float(), vace.float() +for part, sl in (("inactive", slice(0, 16)), ("reactive", slice(16, 32)), ("mask", slice(32, 96))): + d = (a[:, sl] - b[:, sl]).abs() + print(f" {part}: max {d.max().item():.4f} mean {d.mean().item():.5f} ref-frame max {d[:, :, :1].max().item():.4f}") diff --git a/viddub/full_context/check_syncloss.py b/viddub/full_context/check_syncloss.py new file mode 100644 index 0000000000000000000000000000000000000000..28b62e66862819fd9b4216b9e761afe44d938e29 --- /dev/null +++ b/viddub/full_context/check_syncloss.py @@ -0,0 +1,42 @@ +"""Gate for S3: does StableSyncNet on our (non-aligned) crops separate aligned audio from shifted / other-clip audio? + +python check_syncloss.py --n 48 +""" +import argparse + +import numpy as np +import torch + +from .runtime import FPS, load_clip, load_rows +from .sync_loss import N_FRAMES, SyncLoss, melspectrogram, mel_window, mouth_box + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n", type=int, default=48) + args = ap.parse_args() + dev = torch.device("cuda") + loss = SyncLoss().to(dev) + clips = [load_clip(r) for r in load_rows("val")[:args.n]] + mels = [melspectrogram(c["audio"]) for c in clips] + variants = [(0.0, 0.0), (0.1, 0.0), (0.25, 0.0), (0.1, 0.15), (0.25, 0.3)] + for pad_x, pad_top in variants: + pos, shift, other = [], [], [] + for i, c in enumerate(clips): + box = mouth_box(c["mask"], pad_x, pad_top) + video = torch.from_numpy(c["frames"]).to(dev).permute(3, 0, 1, 2).float() / 127.5 - 1 + for s in (1, 33, 64): + v = video[:, s:s + N_FRAMES] + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + pos.append(loss.cos(v, box, mel_window(mels[i], s / FPS)).item()) + s2 = s + 12 if s + 12 + N_FRAMES <= 81 else s - 12 + shift.append(loss.cos(v, box, mel_window(mels[i], s2 / FPS)).item()) + other.append(loss.cos(v, box, mel_window(mels[(i + 1) % len(clips)], s / FPS)).item()) + pos, shift, other = map(np.array, (pos, shift, other)) + print(f"pad_x {pad_x:.2f} pad_top {pad_top:.2f} | cos pos {pos.mean():.3f} shift {shift.mean():.3f} other {other.mean():.3f} | " + f"acc vs shift {(pos > shift).mean():.2f} vs other {(pos > other).mean():.2f} | bce pos {-np.log(np.clip(pos, 1e-6, 1)).mean():.3f}", + flush=True) + + +if __name__ == "__main__": + main() diff --git a/viddub/full_context/color_harmonization.py b/viddub/full_context/color_harmonization.py new file mode 100644 index 0000000000000000000000000000000000000000..6bc5221aad9f34c705e2283ba2a6e0ed4f7c9322 --- /dev/null +++ b/viddub/full_context/color_harmonization.py @@ -0,0 +1,155 @@ +"""Color harmonization for source-video lower-face editing.""" +from __future__ import annotations + +from dataclasses import dataclass + +import cv2 +import numpy as np + + +@dataclass(frozen=True) +class HarmonizationConfig: + feather_px: float = 18.0 + anchor_side_fraction: float = 0.22 + anchor_top_fraction: float = 0.24 + temporal_window: int = 9 + max_lab_delta: tuple[float, float, float] = (18.0, 20.0, 20.0) + spatial_sigma_px: float = 24.0 + luminance_field_strength: float = 0.35 + chroma_field_strength: float = 0.90 + max_spatial_delta: tuple[float, float, float] = (18.0, 30.0, 30.0) + + +def _bbox(mask: np.ndarray) -> tuple[int, int, int, int]: + ys, xs = np.nonzero(mask) + if not len(xs): + raise ValueError("editable mask is empty") + return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1 + + +def skin_anchor_mask(mask: np.ndarray, config: HarmonizationConfig) -> np.ndarray: + """Select cheek and upper-mask anchors while excluding the central mouth.""" + mask = np.asarray(mask, dtype=bool) + left, top, right, bottom = _bbox(mask) + width, height = right - left, bottom - top + side = max(4, round(width * config.anchor_side_fraction)) + upper = max(4, round(height * config.anchor_top_fraction)) + yy, xx = np.ogrid[: mask.shape[0], : mask.shape[1]] + geometry = (xx < left + side) | (xx >= right - side) | (yy < top + upper) + interior = cv2.erode(mask.astype(np.uint8), np.ones((5, 5), np.uint8)) > 0 + anchors = mask & geometry & interior + if anchors.sum() < 64: + raise ValueError("editable mask has too few skin anchors") + return anchors + + +def feather_alpha(mask: np.ndarray, feather_px: float) -> np.ndarray: + """Return an alpha that is zero outside and ramps inward from every edge.""" + binary = np.asarray(mask, dtype=bool).astype(np.uint8) + padded = np.pad(binary, 1) + distance = cv2.distanceTransform(padded, cv2.DIST_L2, 5)[1:-1, 1:-1] + alpha = np.clip(distance / max(float(feather_px), 1.0), 0.0, 1.0) + return alpha * alpha * (3.0 - 2.0 * alpha) + + +def _smooth_deltas(values: np.ndarray, window: int) -> np.ndarray: + window = max(1, int(window) | 1) + radius = window // 2 + if radius == 0: + return values + padded = np.pad(values, ((radius, radius), (0, 0)), mode="edge") + return np.stack( + [np.median(padded[index:index + window], axis=0) + for index in range(len(values))] + ) + + +def estimate_lab_deltas( + generated: np.ndarray, + source: np.ndarray, + anchors: np.ndarray, + config: HarmonizationConfig, +) -> np.ndarray: + """Estimate temporally stable source-minus-generated Lab offsets.""" + if generated.shape != source.shape or generated.ndim != 4 or generated.shape[-1] != 3: + raise ValueError("generated and source must be matching [T,H,W,3] RGB arrays") + raw = [] + for generated_frame, source_frame in zip(generated, source): + generated_lab = cv2.cvtColor( + generated_frame, cv2.COLOR_RGB2LAB).astype(np.float32) + source_lab = cv2.cvtColor( + source_frame, cv2.COLOR_RGB2LAB).astype(np.float32) + valid = anchors & (source_lab[..., 0] > 35) & (source_lab[..., 0] < 235) + if valid.sum() < 64: + valid = anchors + raw.append(np.median(source_lab[valid] - generated_lab[valid], axis=0)) + deltas = _smooth_deltas( + np.asarray(raw, dtype=np.float32), config.temporal_window) + limits = np.asarray(config.max_lab_delta, dtype=np.float32) + return np.clip(deltas, -limits, limits) + + +def harmonize_lower_face( + generated: np.ndarray, + source: np.ndarray, + mask: np.ndarray, + config: HarmonizationConfig | None = None, +) -> tuple[np.ndarray, dict]: + """Correct masked-region color and composite it into exact source context.""" + config = config or HarmonizationConfig() + generated = np.asarray(generated, dtype=np.uint8) + source = np.asarray(source, dtype=np.uint8) + anchors = skin_anchor_mask(mask, config) + alpha = feather_alpha(mask, config.feather_px)[..., None].astype(np.float32) + deltas = estimate_lab_deltas(generated, source, anchors, config) + + output = np.empty_like(generated) + before = [] + after = [] + interior = np.asarray(mask, dtype=bool) & (alpha[..., 0] > 0.95) + for index, (generated_frame, source_frame) in enumerate( + zip(generated, source) + ): + generated_lab = cv2.cvtColor( + generated_frame, cv2.COLOR_RGB2LAB).astype(np.float32) + source_lab = cv2.cvtColor( + source_frame, cv2.COLOR_RGB2LAB).astype(np.float32) + base_lab = generated_lab + deltas[index] + field = cv2.GaussianBlur( + source_lab - base_lab, + (0, 0), + config.spatial_sigma_px, + borderType=cv2.BORDER_REFLECT, + ) + field[..., 0] *= config.luminance_field_strength + field[..., 1:] *= config.chroma_field_strength + limits = np.asarray(config.max_spatial_delta, dtype=np.float32) + field = np.clip(field, -limits, limits) + corrected_lab = np.clip(base_lab + field, 0, 255).astype(np.uint8) + corrected = cv2.cvtColor( + corrected_lab, cv2.COLOR_LAB2RGB).astype(np.float32) + composite = ( + source_frame.astype(np.float32) * (1.0 - alpha) + corrected * alpha + ) + output[index] = np.clip(np.rint(composite), 0, 255).astype(np.uint8) + before.append(np.mean(np.abs( + generated_lab[interior][:, 1:] - source_lab[interior][:, 1:] + ))) + output_lab = cv2.cvtColor( + output[index], cv2.COLOR_RGB2LAB).astype(np.float32) + after.append(np.mean(np.abs( + output_lab[interior][:, 1:] - source_lab[interior][:, 1:] + ))) + + report = { + "method": "low_frequency_lab_field_source_feather", + "feather_px": config.feather_px, + "spatial_sigma_px": config.spatial_sigma_px, + "luminance_field_strength": config.luminance_field_strength, + "chroma_field_strength": config.chroma_field_strength, + "anchor_pixels": int(anchors.sum()), + "median_lab_delta": np.median(deltas, axis=0).round(4).tolist(), + "interior_chroma_mae_before": float(np.mean(before)), + "interior_chroma_mae_after": float(np.mean(after)), + } + return output, report diff --git a/viddub/full_context/inference.py b/viddub/full_context/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..46626ece853829a57f2d7538ac43e0451758154c --- /dev/null +++ b/viddub/full_context/inference.py @@ -0,0 +1,247 @@ +"""Legacy SFT inference utility for component checkpoints. + +This module retains the multi-step teacher/SFT path used during development. +Public two-step VidDub inference must use ``python scripts/infer.py`` so the +serialized release config, DMD sampler, native timing, and overlap contract are +validated together. +""" +import argparse +from pathlib import Path + +import cv2 +import decord +import librosa +import numpy as np +import torch +from safetensors.torch import load_file + +from .audio import DubAudioAdapter +from .color_harmonization import harmonize_lower_face +from .runtime import FPS, PROJECT, SR, ZeroLatents, load_models, sample, write_mp4 + +from data.s01_analyze import detect_landmarks # noqa: E402 +from data.s02_make_clips import CROP_SCALE, LM_LIP_LOW, LM_LIP_UP, NUM_FRAMES, SIZE, crop_box, lower_face_mask # noqa: E402 + + +PAPER_OVERLAP = 17 + + +def windows(n_frames, overlap=PAPER_OVERLAP, require_stride_multiple=False): + """Return bounded source-reanchored windows, shifting the last to cover the tail.""" + stride = NUM_FRAMES - overlap + if n_frames < NUM_FRAMES: + raise ValueError(f"video has {n_frames} frames, need at least {NUM_FRAMES}") + if stride <= 0 or (require_stride_multiple and stride % 4): + raise ValueError( + f"window stride must be positive" + f"{' and a multiple of four' if require_stride_multiple else ''}; " + f"got {stride}") + starts = list(range(0, n_frames - NUM_FRAMES + 1, stride)) + if starts[-1] + NUM_FRAMES < n_frames: + starts.append(n_frames - NUM_FRAMES) + return starts + + +def paste_back(frame, gen, box, alpha): + """Blend the generated crop back into the full frame inside the feathered mouth mask.""" + x0, y0, side = box + g = cv2.resize(gen, (side, side), interpolation=cv2.INTER_CUBIC) + a = cv2.resize(alpha, (side, side), interpolation=cv2.INTER_LINEAR)[..., None] + out = frame.copy() + region = out[y0:y0 + side, x0:x0 + side].astype(np.float32) + out[y0:y0 + side, x0:x0 + side] = (region * (1 - a) + g.astype(np.float32) * a).round().astype(np.uint8) + return out + + +def load_dub_models(adapter_path, dit_path=None, vace_path=None, device=torch.device("cuda")): + pipe, whisper = load_models(device) + if dit_path: + pipe.dit.load_state_dict(load_file(dit_path)) + if vace_path: + pipe.vace.load_state_dict(load_file(vace_path)) + adapter = DubAudioAdapter().to(device) + adapter.load_state_dict(load_file(adapter_path)) + adapter.eval() + return pipe, whisper, adapter, ZeroLatents(pipe) + + +def dub_video( + models, + frames, + audio, + prompt="", + steps=30, + audio_cfg=1.0, + feather=15, + seed=0, + window_box=False, + fps=FPS, + stretch=False, + overlap=PAPER_OVERLAP, + mode="auto", + harmonize=True, + log=print, + sample_fn=None, +): + """frames: uint8 [T,H,W,3] at 24 fps; audio: 16 kHz mono. Returns dubbed frames, same shape. + window_box: crop box and mask from each window's own landmarks (as in training) instead of the whole video. + fps: native frame rate of `frames`. The audio tokens are taken at the native frame times, or (stretch) the audio is + time-stretched (pitch kept) to the training frame rate. + mode: `clip` makes one 81-frame call; `windowed` uses bounded source-reanchored windows; `auto` selects by length. + overlap: frames shared by consecutive windows, linearly cross-faded in crop space (static crop box only). + harmonize: apply the paper's lower-face color correction before window blending.""" + pipe, whisper, adapter, zero_latents = models + n, H, W = len(frames), *frames.shape[1:3] + if mode not in {"auto", "clip", "windowed"}: + raise ValueError(f"unknown inference mode: {mode}") + mode = "clip" if mode == "auto" and n == NUM_FRAMES else ( + "windowed" if mode == "auto" else mode) + if mode == "clip" and n != NUM_FRAMES: + raise ValueError( + f"clip mode requires exactly {NUM_FRAMES} frames; got {n}") + if window_box and mode == "windowed" and overlap: + raise ValueError("overlap blending requires one stable full-video crop") + starts = [0] if mode == "clip" else windows( + n, overlap, require_stride_multiple=sample_fn is not None) + audio = np.pad(audio, (0, max(0, int(n / fps * SR) - len(audio)))) + if stretch and fps != FPS: + audio = librosa.effects.time_stretch(audio, rate=FPS / fps) + t_fps = FPS if stretch else fps + + lm, _ = detect_landmarks(frames) + if not np.isfinite(lm).any(): + # MediaPipe's default detector can miss clear but relatively small + # faces. Retry only complete misses so the standard path and benchmark + # behavior remain unchanged. + lm, _ = detect_landmarks( + frames, + min_detection_confidence=0.05, + min_tracking_confidence=0.1, + ) + if not np.isfinite(lm).any(): + raise RuntimeError("no face landmarks detected in any frame") + log( + "default face detector missed all frames; " + "using low-confidence FaceMesh fallback" + ) + # Tight shots (face*CROP_SCALE > frame): edge-pad so the face/crop ratio matches training. + face = np.nanmedian(np.maximum(np.ptp(lm[:, :, 0], 1), np.ptp(lm[:, :, 1], 1))) + if not np.isfinite(face): + raise RuntimeError("face landmark extent is not finite") + pad = max(0, int(np.ceil((face * CROP_SCALE - min(H, W)) / 2)) + 1) + if pad > 1: + frames_p = np.pad(frames, ((0, 0), (pad, pad), (pad, pad), (0, 0)), mode="edge") + lm = lm + pad + else: + pad, frames_p = 0, frames + ref_idx = int(np.nanargmax(lm[:, LM_LIP_LOW, 1] - lm[:, LM_LIP_UP, 1])) + + def region(lm_w): + box = crop_box(lm_w, H + 2 * pad, W + 2 * pad) + mask = lower_face_mask(lm_w, box) + return box, mask, cv2.GaussianBlur(mask, (0, 0), feather).astype(np.float32) / 255 + + def crop(f, box): + x0, y0, side = box + return cv2.resize(f[y0:y0 + side, x0:x0 + side], (SIZE, SIZE), interpolation=cv2.INTER_AREA) + + # Default: one crop box and one mask for the whole video, keeping the mouth region stable across windows. + box, mask, alpha = region(lm) + out = frames_p.copy() + crop_output = None + written = 0 + if not window_box: + source_crops = np.stack([crop(f, box) for f in frames_p]) + crop_output = np.empty_like(source_crops) + for window_index, start in enumerate(starts): + if window_box: + box, mask, alpha = region(lm[start:start + NUM_FRAMES]) + crops = np.stack([ + crop(f, box) for f in frames_p[start:start + NUM_FRAMES] + ]) + else: + crops = source_crops[start:start + NUM_FRAMES] + clip = {"frames": crops, "mask": mask > 127, "ref": crop(frames_p[ref_idx], box), "prompt": prompt} + sample_start = round(start / t_fps * SR) + sample_count = round(NUM_FRAMES / t_fps * SR) + a = audio[sample_start:sample_start + sample_count] + a = np.pad(a, (0, max(0, sample_count - len(a)))) + sampler = sample if sample_fn is None else sample_fn + window_seed = seed + window_index + gen = sampler( + pipe, whisper, adapter, zero_latents, clip, a, + steps=steps, audio_cfg=audio_cfg, seed=window_seed, fps=t_fps) + color_report = None + if harmonize: + gen, color_report = harmonize_lower_face(gen, crops, mask > 127) + if window_box: + paste_alpha = (mask > 127).astype(np.float32) if harmonize else alpha + for i, g in enumerate(gen): + out[start + i] = paste_back( + frames_p[start + i], g, box, paste_alpha) + written = max(written, start + len(gen)) + else: + end = start + len(gen) + ov = max(0, written - start) + if ov: + weight = ( + np.arange(1, ov + 1, dtype=np.float32) / (ov + 1) + )[:, None, None, None] + crop_output[start:start + ov] = np.rint( + crop_output[start:start + ov].astype(np.float32) * (1 - weight) + + gen[:ov].astype(np.float32) * weight + ).astype(np.uint8) + crop_output[start + ov:end] = gen[ov:] + written = max(written, end) + detail = "" + if color_report is not None: + detail = ( + f", chroma MAE {color_report['interior_chroma_mae_before']:.3f}" + f"->{color_report['interior_chroma_mae_after']:.3f}" + ) + log( + f"window {window_index}: {start}-{start + NUM_FRAMES}, " + f"seed {window_seed}{detail}") + if written != n: + raise RuntimeError(f"window merge wrote {written} frames, expected {n}") + if not window_box: + paste_alpha = (mask > 127).astype(np.float32) if harmonize else alpha + for index, generated in enumerate(crop_output): + out[index] = paste_back( + frames_p[index], generated, box, paste_alpha) + return out[:, pad:pad + H, pad:pad + W] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--video", required=True) + ap.add_argument("--audio", required=True) + ap.add_argument("--adapter", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--dit", default=None) + ap.add_argument("--vace", default=None) + ap.add_argument("--prompt", default="", help="ignored (kept for old command lines): the text context is always empty") + ap.add_argument("--steps", type=int, default=30) + ap.add_argument("--audio_cfg", type=float, default=1.0) + ap.add_argument("--feather", type=int, default=15, help="mask feather radius in crop pixels") + ap.add_argument("--mode", choices=("auto", "clip", "windowed"), default="auto") + ap.add_argument("--overlap", type=int, default=PAPER_OVERLAP, + help="shared frames (paper protocol: 17, giving stride 64)") + ap.add_argument("--no_harmonize", action="store_true", + help="disable paper lower-face color harmonization") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--window_box", action="store_true", help="per-window crop box and mask") + args = ap.parse_args() + + models = load_dub_models(args.adapter, args.dit, args.vace) + frames = decord.VideoReader(args.video, num_threads=2)[:].asnumpy() + audio = librosa.load(args.audio, sr=SR, mono=True)[0] + out = dub_video(models, frames, audio, args.prompt, args.steps, args.audio_cfg, + args.feather, args.seed, args.window_box, overlap=args.overlap, + mode=args.mode, harmonize=not args.no_harmonize) + write_mp4(args.out, out, audio[:int(len(out) / FPS * SR)]) + print("wrote", args.out) + + +if __name__ == "__main__": + main() diff --git a/viddub/full_context/lora.py b/viddub/full_context/lora.py new file mode 100644 index 0000000000000000000000000000000000000000..c19043ae8d750ddbd4137ab42e85471eabefdf03 --- /dev/null +++ b/viddub/full_context/lora.py @@ -0,0 +1,47 @@ +"""Minimal LoRA for the Wan DiT linear layers (attention q/k/v/o + FFN), merged back into plain weights on save.""" +import torch +import torch.nn as nn + +TARGETS = ("self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o", + "cross_attn.q", "cross_attn.k", "cross_attn.v", "cross_attn.o", "ffn.0", "ffn.2") + + +class LoRALinear(nn.Module): + def __init__(self, base: nn.Linear, rank, alpha): + super().__init__() + self.base, self.scale = base, alpha / rank + self.down = nn.Parameter(torch.randn(rank, base.in_features) / rank) + self.up = nn.Parameter(torch.zeros(base.out_features, rank)) + + def forward(self, x): + return self.base(x) + (x @ self.down.t().to(x.dtype)) @ self.up.t().to(x.dtype) * self.scale + + def merged_weight(self): + return self.base.weight + (self.up.float() @ self.down.float() * self.scale).to(self.base.weight.dtype) + + +def inject(dit, rank=64, alpha=64): + """Wrap the target linears of every DiT block; returns the LoRA parameters.""" + params = [] + for block in dit.blocks: + for name in TARGETS: + parent_name, child = name.rsplit(".", 1) + parent = block.get_submodule(parent_name) + lora = LoRALinear(getattr(parent, child), rank, alpha).to(getattr(parent, child).weight.device) + setattr(parent, child, lora) + params += [lora.down, lora.up] + return params + + +def merged_state_dict(dit): + """DiT state dict with the original keys (LoRA folded into the weights).""" + sd = {} + for name, module in dit.named_modules(): + if isinstance(module, LoRALinear): + sd[f"{name}.weight"] = module.merged_weight().detach() + sd[f"{name}.bias"] = module.base.bias.detach() + for k, v in dit.state_dict().items(): + if ".base." in k or k.endswith(".down") or k.endswith(".up"): + continue + sd.setdefault(k, v) + return sd diff --git a/viddub/full_context/repa.py b/viddub/full_context/repa.py new file mode 100644 index 0000000000000000000000000000000000000000..a0967d70c24034284e0b97080020254674fa6533 --- /dev/null +++ b/viddub/full_context/repa.py @@ -0,0 +1,70 @@ +"""Representation alignment of DiT mouth tokens (hidden states after one DiT block) to frozen encoders. + +art: StableSyncNet embedding space. 17 windows of 16 frames (start 4k-3 -> latent frames k..k+3); a projector maps the + concatenated mouth-pooled tokens of those 4 latent frames to 2048-d. Targets: the visual embedding of the ground-truth + window (same audio) or, with `cross`, the audio embedding of the driving audio window (counterfactual steps, no GT + video exists). `temporal` adds a loss on window-to-window differences. +dino: per-token alignment of the mouth-box tokens of every latent frame to DINOv2 patch features of the GT mouth box. +""" +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .runtime import FPS, WEIGHTS +from .sync_loss import mel_window + +N_WIN, WIN_LAT = 17, 4 +DINO_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) +DINO_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + + +def mlp(i, h, o): + return nn.Sequential(nn.Linear(i, h), nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, o)) + + +class Repa(nn.Module): + def __init__(self, kind, dim=1536, temporal=False, cross=False): + super().__init__() + self.kind, self.temporal, self.cross = kind, temporal and kind == "art", cross and kind == "art" + if kind == "art": + self.proj = mlp(WIN_LAT * dim, 2048, 2048) + else: + from transformers import Dinov2Model + self.proj = mlp(dim, 2048, 768) + self.dino = [Dinov2Model.from_pretrained(str(WEIGHTS / "dinov2-base")).eval().requires_grad_(False)] # not a submodule + + def to(self, *a, **kw): + if self.kind == "dino": + self.dino[0].to(*a, **kw) + return super().to(*a, **kw) + + def art_loss(self, h, tok_mask, syncnet, video_gt, box, mel, cross): + """h: [F, 32, 32, C] (latent frames without ref); tok_mask: [32, 32] bool; video_gt: [3, 81, h, w] or None.""" + f = h[:, tok_mask].mean(1) # F, C + x = torch.stack([f[k:k + WIN_LAT].flatten() for k in range(1, N_WIN + 1)]) + p = F.normalize(self.proj(x).float(), dim=-1) + with torch.no_grad(): + if cross: + m = torch.cat([mel_window(mel, (4 * k - 3) / FPS) for k in range(1, N_WIN + 1)]).to(h.device) + t = F.normalize(syncnet.net.audio_encoder(m).flatten(1).float(), dim=-1) + else: + v = torch.cat([syncnet.crop(video_gt[:, 4 * k - 3:4 * k + 13], box) for k in range(1, N_WIN + 1)]) + t = F.normalize(syncnet.net.visual_encoder(v).flatten(1).float(), dim=-1) + loss = (1 - (p * t).sum(-1)).mean() + if self.temporal: + loss = loss + ((p[1:] - p[:-1]) - (t[1:] - t[:-1])).pow(2).sum(-1).mean() + return loss + + def dino_loss(self, h, video_gt, box, tok_box): + """h: [F, 32, 32, C]; video_gt: [3, 81, h, w] in [-1, 1]; box: pixel box in video_gt; tok_box: token box in h.""" + t, b, l, r = box + frames = video_gt[:, ::4, t:b, l:r].transpose(0, 1).float() # F, 3, bh, bw + with torch.no_grad(): + x = F.interpolate((frames + 1) / 2, size=(224, 224), mode="bilinear", align_corners=False) + x = (x - DINO_MEAN.to(x)) / DINO_STD.to(x) + target = F.normalize(self.dino[0](pixel_values=x).last_hidden_state[:, 1:].float(), dim=-1) # F, 256, 768 + tt, tb, tl, tr = tok_box + z = h[:, tt:tb, tl:tr].permute(0, 3, 1, 2).float() + z = F.interpolate(z, size=(16, 16), mode="bilinear", align_corners=False).flatten(2).transpose(1, 2) # F, 256, C + p = F.normalize(self.proj(z.to(h.dtype)).float(), dim=-1) + return (1 - (p * target).sum(-1)).mean() diff --git a/viddub/full_context/runtime.py b/viddub/full_context/runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..9de85aa75d03b15bc71e30a65f48c9af912f0583 --- /dev/null +++ b/viddub/full_context/runtime.py @@ -0,0 +1,232 @@ +"""Shared pieces for Wan-Dub training / validation / inference.""" +import hashlib +import json +import os +import subprocess +from pathlib import Path + +import decord +import imageio_ffmpeg +import numpy as np +import soundfile +import torch +from einops import rearrange +from PIL import Image + +from viddub.foundation.pipelines.wan_video import ModelConfig, WanVideoPipeline, WanVideoUnit_PromptEmbedder, model_fn_wan_video +from .audio import WhisperFeatures + +PROJECT = Path(__file__).resolve().parents[2] +WEIGHTS = Path(os.environ.get("VIDDUB_WEIGHT_ROOT", PROJECT / "weights")) +DATA_ROOT = Path(os.environ.get("VIDDUB_DATA_ROOT", PROJECT / "data")) +CLIPS = DATA_ROOT / "clips" +VACE_DIR = WEIGHTS / "Wan2.1-VACE-1.3B" +SR = 16000 +FPS = 24 +TEXT_LEN = 512 +MANIFEST = DATA_ROOT / "manifest" +SOURCES = {"h3": MANIFEST, "hallo3": MANIFEST / "hallo3", "ls_real": MANIFEST / "ls_real", "astri": MANIFEST / "astri"} + + +def load_rows(split, sources=("h3",)): + """s02 rows of `sources`, each tagged with source / domain / group. val = md5(group) % 100 == 0, so no source video + (HDTF person, VoxCeleb2 speaker, YouTube id) straddles train/val; groups shared with the public benches are dropped + from both. Optional group maps are stored with the data manifests.""" + exclude_path = MANIFEST / "bench_exclude.json" + excluded = ({g for v in json.load(open(exclude_path)).values() for g in v} + if exclude_path.is_file() else set()) + out = [] + for src in sources: + groups = json.load(open(SOURCES[src] / "groups.json")) if src != "h3" else None + for p in sorted((SOURCES[src] / "s02").glob("shard_*.jsonl")): + for l in open(p): + if not l.strip(): + continue + r = json.loads(l) + if "error" in r: + continue + r["source"] = src + r["domain"], r["group"] = groups[r["id"]] if groups else ("H3", r["id"]) + if r["group"] not in excluded and (int(hashlib.md5(r["group"].encode()).hexdigest(), 16) % 100 == 0) == (split == "val"): + out.append(r) + return sorted(out, key=lambda r: r["clip_id"]) + + +def cached_clip_ids(sources=("h3",)): + ids = set() + for src in sources: + for p in (SOURCES[src] / "s03").glob("shard_*.txt"): + ids.update(p.read_text().split()) + return ids + + +def load_clip(row, num_frames=81): + d = CLIPS / row["dir"] + frames = decord.VideoReader(str(d / "video.mp4"), num_threads=1).get_batch(list(range(num_frames))).asnumpy() + audio, _ = soundfile.read(d / "audio.wav", dtype="float32") + return {"frames": frames, "mask": np.array(Image.open(d / "mask.png")) > 127, "ref": np.array(Image.open(d / "ref.png").convert("RGB")), + "audio": audio[:round(num_frames / FPS * SR)], "prompt": row["prompt"], "clip_id": row["clip_id"]} + + +class CachedClipDataset(torch.utils.data.Dataset): + """Precomputed latents (s03) + mask + audio.""" + + def __init__(self, rows, num_frames, latents_name="latents.pt"): + self.rows, self.num_frames, self.latents_name = rows, num_frames, latents_name + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + d = CLIPS / self.rows[i]["dir"] + audio, _ = soundfile.read(d / "audio.wav", dtype="float32") + return {"latents": torch.load(d / self.latents_name), "mask": np.array(Image.open(d / "mask.png")) > 127, + "audio": audio[:round(self.num_frames / FPS * SR)], "clip_id": self.rows[i]["clip_id"]} + + +def load_models(device): + pipe = WanVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, device=device, + model_configs=[ModelConfig(path=str(VACE_DIR / "diffusion_pytorch_model.safetensors")), + ModelConfig(path=str(VACE_DIR / "models_t5_umt5-xxl-enc-bf16.pth")), + ModelConfig(path=str(VACE_DIR / "Wan2.1_VAE.pth"))], + tokenizer_config=ModelConfig(path=str(VACE_DIR / "google/umt5-xxl")), + ) + for m in (pipe.text_encoder, pipe.dit, pipe.vace, pipe.vae): + m.requires_grad_(False).eval() + whisper = WhisperFeatures(str(WEIGHTS / "whisper-small")).to(device, torch.bfloat16) + return pipe, whisper + + +def to_video_tensor(frames, pipe): + """uint8 [T,H,W,3] -> [1,3,T,H,W] in [-1,1].""" + return (torch.from_numpy(frames).to(pipe.device, pipe.torch_dtype).permute(3, 0, 1, 2)[None] / 127.5 - 1) + + +@torch.no_grad() +def encode_text(pipe, prompt): + """T5 context trimmed to its real length: [L, 4096].""" + emb = WanVideoUnit_PromptEmbedder().encode_prompt(pipe, prompt)[0] + n = int((emb.abs().sum(-1) > 0).sum()) + return emb[:max(n, 1)] + + +@torch.no_grad() +def compute_latents(pipe, clip): + """VAE latents of gt video, masked (inactive) video, reference frame; plus T5 context. Matches WanVideoUnit_VACE.""" + video = to_video_tensor(clip["frames"], pipe) + keep = torch.from_numpy(~clip["mask"]).to(video)[None, None, None] + return {"x0": pipe.vae.encode(video, device=pipe.device)[0].cpu(), + "inactive": pipe.vae.encode(video * keep, device=pipe.device)[0].cpu(), + "ref": pipe.vae.encode(to_video_tensor(clip["ref"][None], pipe), device=pipe.device)[0].cpu(), + "context": encode_text(pipe, clip["prompt"]).cpu()} + + +class ZeroLatents: + """VAE latent of an all-zero video (= VACE 'reactive' part, since the masked region is blanked).""" + + def __init__(self, pipe): + self.pipe, self.cache, self._empty_context = pipe, {}, None + + def empty_context(self): + """T5 context of "": the only text condition (dubbing is audio-driven; inference never has a prompt).""" + if self._empty_context is None: + cached = getattr(self.pipe, "_wan_dub_empty_context", None) + self._empty_context = (cached.to(self.pipe.device, self.pipe.torch_dtype) + if cached is not None else encode_text(self.pipe, "")) + return self._empty_context + + def __call__(self, num_frames, h, w): + key = (num_frames, h, w) + if key not in self.cache: + z = torch.zeros(1, 3, num_frames, h, w, device=self.pipe.device, dtype=self.pipe.torch_dtype) + self.cache[key] = self.pipe.vae.encode(z, device=self.pipe.device)[0] + return self.cache[key] + + +def build_inputs(lat, mask, num_frames, zero_latents, drop_ref=False, context=None): + """Cached latents -> (x0 [1,16,F,h,w], vace_context [1,96,F,h,w], context [1,512,4096]) as WanVideoUnit_VACE would. + context=None -> empty-prompt context (the cached caption context lat["context"] is not used unless passed explicitly).""" + device, dtype = zero_latents.pipe.device, zero_latents.pipe.torch_dtype + n_lat = (num_frames - 1) // 4 + 1 + x0 = lat["x0"][:, :n_lat].to(device, dtype) + H, W = mask.shape + vace = torch.cat([lat["inactive"][:, :n_lat].to(device, dtype), zero_latents(num_frames, H, W)[:, :n_lat]], 0) + m = rearrange(torch.from_numpy(mask).to(device, dtype), "(H P) (W Q) -> (P Q) 1 H W", P=8, Q=8).expand(-1, n_lat, -1, -1) + if not drop_ref: + ref = lat["ref"].to(device, dtype) + x0 = torch.cat([ref, x0], 1) + vace = torch.cat([torch.cat([ref, torch.zeros_like(ref)], 0), vace], 1) + m = torch.cat([torch.zeros_like(m[:, :1]), m], 1) + ctx = zero_latents.empty_context() if context is None else context + ctx = torch.cat([ctx, ctx.new_zeros(TEXT_LEN - len(ctx), ctx.shape[1])]).to(device, dtype)[None] + return x0[None], torch.cat([vace, m], 0)[None], ctx + + +def predict(pipe, adapter, latents, timestep, context, vace_context, audio_tokens, grad_ckpt=False): + return model_fn_wan_video(dit=pipe.dit, vace=pipe.vace, latents=latents, timestep=timestep, context=context, + vace_context=vace_context, vace_scale=1, audio_adapter=adapter, audio_tokens=audio_tokens, + use_gradient_checkpointing=grad_ckpt) + + +def mask_weight(vace_context, lam): + """Latent-resolution mouth mask from VACE context (channels 32: are the 8x8-unshuffled pixel mask).""" + return 1 + lam * vace_context[:, 32:].amax(1, keepdim=True).float() + + +@torch.no_grad() +def sample(pipe, whisper, adapter, zero_latents, clip, audio, steps=30, audio_cfg=1.0, seed=0, lat=None, drop_ref=False, fps=FPS): + """Returns uint8 frames [T,H,W,3] generated for `clip` (masked, with ref) driven by `audio`.""" + device, dtype = pipe.device, pipe.torch_dtype + n = len(clip["frames"]) + n_ref = 0 if drop_ref else 1 + lat = compute_latents(pipe, clip) if lat is None else lat + x0, vace_context, context = build_inputs(lat, clip["mask"], n, zero_latents, drop_ref=drop_ref) + hs = whisper([audio], device, dtype) + latents = torch.randn(x0.shape, generator=torch.Generator().manual_seed(seed)).to(device, dtype) + pipe.scheduler.set_timesteps(steps, shift=5.0) + with torch.autocast("cuda", dtype=torch.bfloat16): + tok = adapter.tokens(hs, n, n_ref, torch.tensor([False]), fps) + tok_null = adapter.tokens(hs, n, n_ref, torch.tensor([True]), fps) + for t in pipe.scheduler.timesteps: + tt = t.unsqueeze(0).to(dtype=dtype, device=device) + pred = predict(pipe, adapter, latents, tt, context, vace_context, tok) + if audio_cfg != 1.0: + pred_u = predict(pipe, adapter, latents, tt, context, vace_context, tok_null) + pred = pred_u + audio_cfg * (pred - pred_u) + latents = pipe.scheduler.step(pred, t, latents) + video = pipe.vae.decode(latents[:, :, n_ref:].to(dtype), device=device) + pipe.scheduler.set_timesteps(1000, training=True) + return ((video[0].permute(1, 2, 3, 0).float().clamp(-1, 1) + 1) * 127.5).round().byte().cpu().numpy() + + +def write_mp4(path, frames, audio=None, fps=FPS): + Path(path).parent.mkdir(parents=True, exist_ok=True) + h, w = frames.shape[1:3] + cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", + "-r", str(fps), "-i", "-"] + tmp_wav = None + if audio is not None: + tmp_wav = str(path) + ".wav" + soundfile.write(tmp_wav, audio, SR) + # 48k stereo: 16k mono AAC plays silent in many embedded players + cmd += ["-i", tmp_wav, "-c:a", "aac", "-ar", "48000", "-ac", "2", "-shortest"] + cmd += ["-c:v", "libx264", "-crf", "18", "-pix_fmt", "yuv420p", "-movflags", "+faststart", str(path)] + subprocess.run(cmd, input=np.ascontiguousarray(frames).tobytes(), check=True) + if tmp_wav: + Path(tmp_wav).unlink() + + +def syncnet_scorer(device): + from data.s01_analyze import detect_landmarks, load_syncnet, syncnet_crops, syncnet_score + net = load_syncnet(device) + + def score(frames, audio): + lm, nfaces = detect_landmarks(frames) + if (nfaces > 0).mean() < 0.5: + return None + crops = syncnet_crops(frames, lm) + idx25 = (np.arange(int(len(frames) * 25 / FPS)) * FPS / 25).astype(int) + off, dist, conf = syncnet_score(net, crops[idx25], (audio * 32767).astype(np.int16), device) + return {"offset": off, "dist": dist, "conf": conf} + return score diff --git a/viddub/full_context/sync_loss.py b/viddub/full_context/sync_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..4bfcc9ca49f139ddace763af00869c45627be0e9 --- /dev/null +++ b/viddub/full_context/sync_loss.py @@ -0,0 +1,74 @@ +"""LatentSync-1.6 StableSyncNet as a differentiable lip-sync loss on our 512 face crops. + +StableSyncNet sees the lower half of a 256x256 affine-aligned face (128x256), 16 frames @25fps, and 52 mel columns (80 Hz). +Our crops are not affine-aligned; the lower half is approximated by the clip's mouth mask box (nose bottom -> below chin, +cheek to cheek), widened by `pad_x` on each side. 16 frames @24fps span 0.667s, so the matching 53.3 mel columns are +resampled to 52. +""" +import os +import sys +from pathlib import Path + +import librosa +import numpy as np +import torch +import torch.nn.functional as F +import yaml +from scipy import signal + +from .runtime import FPS, SR, WEIGHTS + +LS_ROOT = Path(os.environ.get("VIDDUB_LATENTSYNC_ROOT", WEIGHTS / "LatentSync")) +N_FRAMES, MEL_HZ, MEL_WIN = 16, 80, 52 +_MEL_BASIS = librosa.filters.mel(sr=SR, n_fft=800, n_mels=80, fmin=55, fmax=7600) + + +def melspectrogram(wav): + """Same as latentsync/utils/audio.py with configs/audio.yaml: [80, T] at 80 Hz, values in [-4, 4].""" + d = librosa.stft(y=signal.lfilter([1, -0.97], [1], wav), n_fft=800, hop_length=200, win_length=800) + s = 20 * np.log10(np.maximum(1e-5, _MEL_BASIS @ np.abs(d))) - 20 + return np.clip(8 * (s + 100) / 100 - 4, -4, 4).astype(np.float32) + + +def mel_window(mel, t0): + """Mel columns for the 16 video frames starting at time t0 (s), resampled to MEL_WIN: [1, 1, 80, 52].""" + a = int(round(t0 * MEL_HZ)) + b = a + int(round(N_FRAMES / FPS * MEL_HZ)) + m = torch.from_numpy(mel[:, a:b]) + if m.shape[1] < b - a: + m = F.pad(m, (0, b - a - m.shape[1]), value=-4.0) + return F.interpolate(m[None, None], size=(80, MEL_WIN), mode="bilinear", align_corners=False) + + +def mouth_box(mask, pad_x=0.1, pad_top=0.0): + """Pixel box (t, b, l, r) of the mouth mask, widened; mask: bool [H, W].""" + ys, xs = np.nonzero(mask) + t, b, l, r = ys.min(), ys.max() + 1, xs.min(), xs.max() + 1 + w, h = r - l, b - t + H, W = mask.shape + return max(0, int(t - pad_top * h)), b, max(0, int(l - pad_x * w)), min(W, int(r + pad_x * w)) + + +class SyncLoss(torch.nn.Module): + def __init__(self): + super().__init__() + sys.path.insert(0, str(LS_ROOT)) + from latentsync.models.stable_syncnet import StableSyncNet + cfg = yaml.safe_load(open(LS_ROOT / "configs/syncnet/syncnet_16_pixel_attn.yaml"))["model"] + self.net = StableSyncNet(cfg, gradient_checkpointing=True) + self.net.load_state_dict(torch.load(LS_ROOT / "checkpoints/stable_syncnet.pt", map_location="cpu", weights_only=True)["state_dict"]) + self.net.requires_grad_(False).eval() + + def crop(self, video, box): + """video: [3, 16, H, W] in [-1, 1] -> [1, 48, 128, 256] (frame-major, like LatentSync).""" + t, b, l, r = box + x = F.interpolate(video[:, :, t:b, l:r].transpose(0, 1).float(), size=(128, 256), mode="bilinear", align_corners=False) + return x.reshape(1, -1, 128, 256) + + def cos(self, video, box, mel): + v, a = self.net(self.crop(video, box), mel.to(video.device)) + return F.cosine_similarity(v.float(), a.float()) + + def forward(self, video, box, mel): + """BCE(cos, 1) as in LatentSync's train_unet.py.""" + return -torch.log(self.cos(video, box, mel).clamp(1e-6, 1.0)).mean() diff --git a/viddub/full_context/train.py b/viddub/full_context/train.py new file mode 100644 index 0000000000000000000000000000000000000000..07dc5434c3ab1e9bff3c6019596798a773a1b81f --- /dev/null +++ b/viddub/full_context/train.py @@ -0,0 +1,545 @@ +"""Wan-Dub training (accelerate DDP). Stage s1: train audio adapter only; stage s2: + VACE + DiT.""" +import argparse +import datetime +import itertools +import json +import math +import os +import random +import time +from pathlib import Path + +import accelerate +import numpy as np +import torch +import torch.utils.checkpoint +from safetensors.torch import load_file, save_file +from torch.utils.tensorboard import SummaryWriter + +from .audio import DubAudioAdapter +from .runtime import (FPS, PROJECT, CachedClipDataset, ZeroLatents, build_inputs, cached_clip_ids, load_clip, load_models, load_rows, + mask_weight, predict, sample, syncnet_scorer, write_mp4) +from . import lora +from .repa import Repa +from .sync_loss import N_FRAMES, SyncLoss, mel_window, melspectrogram, mouth_box + + +def first_item(batch): + return batch[0] + + +class DomainBalancedSampler(torch.utils.data.Sampler): + """Endless indices: domain ~ weights, then a uniform group (source video / speaker) of it, then a uniform clip of + that group, so neither big domains nor clip-rich videos dominate. Per-rank seed (not sharded by accelerate).""" + + def __init__(self, rows, weights, seed): + self.groups = {} + for i, r in enumerate(rows): + self.groups.setdefault(r["domain"], {}).setdefault(r["group"], []).append(i) + self.domains = sorted(self.groups) + self.groups = {d: list(self.groups[d].values()) for d in self.domains} + self.weights = [weights.get(d, 1.0) for d in self.domains] + self.seed = seed + + def __iter__(self): + rng = random.Random(self.seed) + while True: + g = rng.choice(self.groups[rng.choices(self.domains, self.weights)[0]]) + yield rng.choice(g) + + +class EpochSampler(torch.utils.data.Sampler): + """Endless epochs over every clip (a domain's clips listed repeats[domain] times), one shuffle per epoch shared by + all ranks, each rank taking its contiguous slice; `start` = samples this rank already consumed (resume).""" + + def __init__(self, rows, repeats, rank, world, start=0): + self.idx = [i for i, r in enumerate(rows) for _ in range(int(repeats.get(r["domain"], 1)))] + self.rank, self.per, self.start = rank, len(self.idx) // world, start + + def __iter__(self): + k = self.start + while True: + ep, off = divmod(k, self.per) + perm = self.idx[:] + random.Random(ep).shuffle(perm) + yield from perm[self.rank * self.per + off:(self.rank + 1) * self.per] + k = (ep + 1) * self.per + + +class MixSampler(torch.utils.data.Sampler): + """Global sample g: domain ~ shares, then the next clip of that domain's current pass, a weighted random permutation + (keys u^(1/w): every clip once per pass, high-weight clips earlier). All ranks build the same stream and rank r takes + g = r, r + world, ...; `start` = samples this rank already consumed (resume).""" + + def __init__(self, rows, shares, clip_w, rank, world, start=0, seed=0): + by = {} + for i, r in enumerate(rows): + by.setdefault(r["domain"], []).append(i) + self.domains = [d for d in sorted(by) if shares.get(d, 1.0) > 0] + self.idx = [np.array(by[d]) for d in self.domains] + self.w = [np.array([clip_w[i] for i in by[d]]) for d in self.domains] + p = np.array([shares.get(d, 1.0) for d in self.domains], dtype=np.float64) + self.p, self.rank, self.world, self.start, self.seed = p / p.sum(), rank, world, start, seed + + def _perm(self, k, n_pass): + rng = np.random.default_rng([self.seed, k, n_pass]) + return self.idx[k][np.argsort(-rng.random(len(self.idx[k])) ** (1 / self.w[k]))] + + def __iter__(self): + perms = [self._perm(k, 0) for k in range(len(self.domains))] + cur, passes, g, first = [0] * len(self.domains), [0] * len(self.domains), 0, self.start * self.world + self.rank + for block in itertools.count(): + for k in np.random.default_rng([self.seed, 1 << 20, block]).choice(len(self.domains), 1 << 16, p=self.p): + if cur[k] == len(perms[k]): + passes[k] += 1 + perms[k], cur[k] = self._perm(k, passes[k]), 0 + i = perms[k][cur[k]] + cur[k] += 1 + if g >= first and g % self.world == self.rank: + yield int(i) + g += 1 + + +def domain_stats(rows, weights): + by = {} + for r in rows: + d = by.setdefault(r["domain"], {"clips": 0, "groups": set()}) + d["clips"] += 1 + d["groups"].add(r["group"]) + tot = sum(weights.get(d, 1.0) for d in by) + return {d: {"clips": v["clips"], "groups": len(v["groups"]), "share": weights.get(d, 1.0) / tot} for d, v in sorted(by.items())} + + +def interleave_domains(rows, n): + """First n rows taking one per domain in turn (clip_id order within a domain).""" + by = {} + for r in rows: + by.setdefault(r["domain"], []).append(r) + queues, out = [by[d] for d in sorted(by)], [] + for k in range(max(map(len, queues), default=0)): + out += [q[k] for q in queues if k < len(q)] + return out[:n] + + +def latent_box(lat, box, max_h=36, max_w=48): + """Latent window (lt, lb, ll, lr) around a pixel box (+2 latent px, at most max_h x max_w around its centre) and the + pixel box clipped to it.""" + t, b, l, r = box + lt, lb = max(0, t // 8 - 2), min(lat.shape[3], -(-b // 8) + 2) + ll, lr = max(0, l // 8 - 2), min(lat.shape[4], -(-r // 8) + 2) + 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), (max(t, 8 * lt), min(b, 8 * lb), max(l, 8 * ll), min(r, 8 * lr)) + + +@torch.no_grad() +def decode_box(pipe, lat, box): + """Decode only the latent window around `box`: ([3, T, h, w] in [-1, 1], box local to it, clipped global box).""" + (lt, lb, ll, lr), (t, b, l, r) = latent_box(lat, box) + video = pipe.vae.model.decode(lat[:, :, :, lt:lb, ll:lr].to(pipe.torch_dtype), pipe.vae.scale)[0] + pipe.vae.model.clear_cache() + return video, (t - 8 * lt, b - 8 * lt, l - 8 * ll, r - 8 * ll), (t, b, l, r) + + +def sync_term(pipe, syncnet, lat, box, mel, lat_gt=None, lpips_net=None, max_h=36, max_w=48): + """SyncNet loss on the first 16 frames decoded from x0_hat latents [1,16,5,h,w], plus (if lat_gt is given) the L1 and + (if lpips_net is given) the LPIPS to the decoded ground truth inside the mouth box. Only the mouth box (+2 latent px, at most max_h x max_w latent px around + its centre) is decoded: bounds decoder memory.""" + (lt, lb, ll, lr), (t, b, l, r) = latent_box(lat, box, max_h, max_w) + decode = lambda z: pipe.vae.model.decode(z, pipe.vae.scale) + video = torch.utils.checkpoint.checkpoint(decode, lat[:, :, :, lt:lb, ll:lr].to(pipe.torch_dtype), use_reentrant=False)[0] # [3, 17, h, w] + t, b, l, r = t - 8 * lt, b - 8 * lt, l - 8 * ll, r - 8 * ll + l_sync = syncnet(video[:, 1:1 + N_FRAMES], (t, b, l, r), mel) + if lat_gt is None: + return l_sync, None, None + with torch.no_grad(): + gt = decode(lat_gt[:, :, :, lt:lb, ll:lr].to(pipe.torch_dtype))[0] + pipe.vae.model.clear_cache() + pred, gt = video[:, 1:1 + N_FRAMES, t:b, l:r].float(), gt[:, 1:1 + N_FRAMES, t:b, l:r].float() + l_lpips = None + if lpips_net is not None: + with torch.autocast("cuda", enabled=False): + l_lpips = lpips_net(pred.clamp(-1, 1).transpose(0, 1), gt.clamp(-1, 1).transpose(0, 1)).mean() + return l_sync, (pred - gt).abs().mean(), l_lpips + + +class Trainable(torch.nn.Module): + """Single DDP entry point: frozen whisper hidden states + DiffSynth conditions -> weighted flow-matching loss (+ SyncNet loss).""" + + def __init__(self, pipe, adapter, syncnet=None, repa=None, repa_layer=8, lpips_net=None): + super().__init__() + self.adapter, self.dit, self.vace, self.repa = adapter, pipe.dit, pipe.vace, repa + self.pipe, self.syncnet, self.lpips = [pipe], [syncnet], [lpips_net] # not submodules + self.hidden = [None] + if repa is not None: + adapter.blocks[repa_layer].register_forward_hook(lambda m, i, o: self.hidden.__setitem__(0, o)) + + def forward(self, x0, context, vace_context, whisper_hs, num_frames, n_ref, drop_audio, lam, sync=None): + """sync: None or {box, mel, weight, sigma_max, cross}. cross = audio from another clip: no MSE inside the mouth, low-noise t only.""" + pipe = self.pipe[0] + tok = self.adapter.tokens(whisper_hs, num_frames, n_ref, drop_audio) + ts = pipe.scheduler.timesteps + if sync and sync["cross"]: + ts = ts[pipe.scheduler.sigmas <= sync["sigma_max"]] + timestep = ts[torch.randint(0, len(ts), (1,))].to(x0.device, x0.dtype) + sigma = float(pipe.scheduler.sigmas[torch.argmin((pipe.scheduler.timesteps - timestep.cpu().float()).abs())]) + noise = torch.randn_like(x0) + xt = pipe.scheduler.add_noise(x0, noise, timestep) + target = pipe.scheduler.training_target(x0, noise, timestep) + pred = predict(pipe, self.adapter, xt, timestep, context, vace_context, tok, grad_ckpt=True) + w = 1 - vace_context[:, 32:].amax(1, keepdim=True).float() if sync and sync["cross"] else mask_weight(vace_context, lam) + loss = (w * (pred.float() - target.float()) ** 2).mean() / w.mean() * pipe.scheduler.training_weight(timestep) + l_sync, l_pix, l_lpips = None, None, None + if sync and sigma <= sync["sigma_max"]: + x0_hat = xt.float() - sigma * pred.float() + pix = (sync["pix_weight"] > 0 or sync["lpips_weight"] > 0) and not sync["cross"] + l_sync, l_pix, l_lpips = sync_term(pipe, self.syncnet[0], x0_hat[:, :, n_ref:n_ref + 5], sync["box"], sync["mel"], + x0[:, :, n_ref:n_ref + 5] if pix else None, + self.lpips[0] if sync["lpips_weight"] > 0 else None) + loss = loss + sync["weight"] * l_sync + if l_pix is not None: + loss = loss + sync["pix_weight"] * l_pix + if l_lpips is not None: + loss = loss + sync["lpips_weight"] * l_lpips + l_repa = None + if self.repa is not None and sync and (not sync["cross"] or any(r.cross for r in self.repa)): + l_repa = self.repa_term(x0, vace_context, n_ref, sync) + loss = loss + sync["repa_weight"] * l_repa + return loss, l_sync, l_pix, l_repa, l_lpips + + def repa_term(self, x0, vace_context, n_ref, sync): + F_, hh, ww = x0.shape[2], x0.shape[3] // 2, x0.shape[4] // 2 + h = self.hidden[0][0].view(F_, hh, ww, -1)[n_ref:] + assert h.requires_grad, "hooked hidden state is detached" + if sync["repa_detach"]: + h = h.detach() + video, box, (t, b, l, r) = (None, None, (0, 0, 0, 0)) if sync["cross"] else decode_box(self.pipe[0], x0[:, :, n_ref:], sync["box"]) + loss = 0 + for repa in self.repa: + if sync["cross"] and not repa.cross: + continue + if repa.kind == "art": + tok_mask = torch.nn.functional.max_pool2d(vace_context[0, 32:, -1].amax(0)[None].float(), 2)[0] > 0 + loss = loss + repa.art_loss(h, tok_mask, self.syncnet[0], video, box, sync["mel_full"], sync["cross"]) + else: + loss = loss + repa.dino_loss(h, video, box, (t // 16, -(-b // 16), l // 16, -(-r // 16))) + return loss + + +def validate(args, accel, pipe, whisper, adapter, zero_latents, val_clips, scorer, step, writer): + adapter.eval() + res = {"same": [], "cross": [], "mouth_l1": []} + for i, clip in enumerate(val_clips): + other = val_clips[(i + 1) % len(val_clips)] + for kind, audio in (("same", clip["audio"]), ("cross", other["audio"])): + frames = sample(pipe, whisper, adapter, zero_latents, clip, audio, steps=args.val_steps_sampling, + audio_cfg=args.val_audio_cfg, seed=i, lat=clip["lat"]) + s = scorer(frames, audio) + res[kind].append(s["conf"] if s else 0.0) + res.setdefault(f"{kind}/{clip['domain']}", []).append(res[kind][-1]) + if kind == "same": + m = clip["mask"] + res["mouth_l1"].append(float(np.abs(frames[:, m].astype(np.float32) - clip["frames"][:, m]).mean())) + if i < args.val_save: + write_mp4(Path(args.output) / "val" / f"step{step:06d}" / f"{clip['clip_id']}_{kind}.mp4", + np.concatenate([clip["frames"], frames], 2), audio) + out = {k: float(np.mean(v)) for k, v in res.items()} + per_domain = " ".join(f"{k[5:]} {v:.2f}" for k, v in out.items() if k.startswith("same/")) + accel.print(f"[val step {step}] sync_conf same {out['same']:.3f} cross {out['cross']:.3f} mouth_l1 {out['mouth_l1']:.2f} | same: {per_domain}", flush=True) + for k, v in out.items(): + writer.add_scalar(f"val/{k}", v, step) + with open(Path(args.output) / "val_log.jsonl", "a") as f: + f.write(json.dumps({"step": step, **out}) + "\n") + adapter.train() + + +def save_weights(args, raw, out, step): + save_file({k: v.detach().float().cpu().contiguous() for k, v in raw.adapter.state_dict().items()}, + str(out / f"adapter_step{step:06d}.safetensors")) + if args.stage == "s2": + dit_sd = lora.merged_state_dict(raw.dit) if args.dit_lora_rank else raw.dit.state_dict() + save_file({k: v.detach().bfloat16().cpu().contiguous() for k, v in dit_sd.items()}, str(out / f"dit_step{step:06d}.safetensors")) + save_file({k: v.detach().bfloat16().cpu().contiguous() for k, v in raw.vace.state_dict().items()}, + str(out / f"vace_step{step:06d}.safetensors")) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--output", required=True) + ap.add_argument("--stage", choices=["s1", "s2"], default="s1") + ap.add_argument("--init_adapter", default=None, help="adapter safetensors to start from (s2)") + ap.add_argument("--init_dit", default=None, help="DiT safetensors (continue an s2 run in a new dir)") + ap.add_argument("--init_vace", default=None, help="VACE safetensors (continue an s2 run in a new dir)") + ap.add_argument("--num_frames", type=int, default=81) + ap.add_argument("--max_train_clips", type=int, default=None) + ap.add_argument("--lr", type=float, default=1e-4) + ap.add_argument("--lr_base", type=float, default=1e-5, help="lr for DiT/VACE in s2") + ap.add_argument("--warmup", type=int, default=200) + ap.add_argument("--max_steps", type=int, default=10000) + ap.add_argument("--grad_accum", type=int, default=1) + ap.add_argument("--mask_lambda", type=float, default=2.0) + ap.add_argument("--p_drop_text", type=float, default=1.0, + help="prob. of the empty-prompt context; 1 = never condition on the cached caption (inference has no text)") + ap.add_argument("--data", default="h3", help="comma list of sources: h3, hallo3, ls_real, astri") + ap.add_argument("--domain_weights", default="", + help="e.g. H3=1,TalkVid=2 (domains: H3 hallo3 HDTF TalkVid SpeakerVid VoxCeleb2; default 1 each). " + "Multi-source runs sample domain -> source video -> clip") + ap.add_argument("--sampling", choices=["balanced", "epoch", "mix"], default="balanced", + help="epoch: every clip once per epoch, --domain_weights are integer repeat factors (e.g. VoxCeleb2=8); " + "mix: --domain_weights are sampling shares, clips of a domain cycle without replacement, ordered by --clip_weights") + ap.add_argument("--clip_weights", default="", help="mix: json {clip_id: weight}; clips not listed are dropped") + ap.add_argument("--p_drop_audio", type=float, default=0.1) + ap.add_argument("--p_drop_ref", type=float, default=0.1) + ap.add_argument("--log_every", type=int, default=20) + ap.add_argument("--save_every", type=int, default=1000) + ap.add_argument("--val_every", type=int, default=1000) + ap.add_argument("--val_clips", type=int, default=8) + ap.add_argument("--val_save", type=int, default=4) + ap.add_argument("--val_steps_sampling", type=int, default=30) + ap.add_argument("--val_audio_cfg", type=float, default=1.0) + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--latents_name", default="latents.pt", help="filename under each clip dir (s03 cache)") + ap.add_argument("--sync_weight", type=float, default=0.0, help="StableSyncNet loss weight (0 = off)") + ap.add_argument("--sync_sigma_max", type=float, default=0.8, help="SyncNet loss only when sigma <= this") + ap.add_argument("--p_cross", type=float, default=0.0, help="prob. of driving with another clip's audio (SyncNet + non-mouth MSE)") + ap.add_argument("--pix_weight", type=float, default=0.0, help="L1 on the decoded mouth box vs decoded GT (same-audio sync steps)") + ap.add_argument("--repa", default="none", help="mouth-token representation alignment: none or a comma list of art, dino") + ap.add_argument("--repa_cross", action="store_true", help="art: align cross-audio steps to the driving audio's SyncNet embedding") + ap.add_argument("--repa_temporal", action="store_true", help="art: also align window-to-window differences") + ap.add_argument("--repa_layer", type=int, default=8) + ap.add_argument("--repa_weight", type=float, default=0.1) + ap.add_argument("--repa_lr", type=float, default=1e-4) + ap.add_argument("--repa_warmup", type=int, default=50, help="steps that only train the projector (hidden state detached)") + ap.add_argument("--lpips_weight", type=float, default=0.0, help="LPIPS(vgg) on the decoded mouth box vs decoded GT (same-audio aux steps)") + ap.add_argument("--dit_lora_rank", type=int, default=0, help="s2 + --freeze_dit: train a LoRA of this rank on the DiT (merged on save)") + ap.add_argument("--lora_lr", type=float, default=1e-4) + ap.add_argument("--lr_schedule", choices=["const", "cosine"], default="const", help="cosine: decay to 0.1x lr at max_steps") + ap.add_argument("--ema", type=float, default=0.0, help="EMA decay of the trained weights (0 = off); saved to _ema/") + ap.add_argument("--skip_val0", action="store_true") + ap.add_argument("--freeze_dit", action="store_true", help="s2: keep the DiT frozen (bf16), train adapter + VACE only") + args = ap.parse_args() + + accel = accelerate.Accelerator(gradient_accumulation_steps=args.grad_accum, step_scheduler_with_optimizer=False, + kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=True, gradient_as_bucket_view=True), + accelerate.InitProcessGroupKwargs(timeout=datetime.timedelta(hours=2))]) + out = Path(args.output) + out.mkdir(parents=True, exist_ok=True) + torch.manual_seed(1234 + accel.process_index) + random.seed(1234 + accel.process_index) + + pipe, whisper = load_models(accel.device) + pipe.scheduler.set_timesteps(1000, training=True) + adapter = DubAudioAdapter().to(accel.device) + if args.init_adapter: + adapter.load_state_dict(load_file(args.init_adapter)) + groups = [{"params": list(adapter.parameters()), "lr": args.lr}] + if args.stage == "s2": + pipe.vace.requires_grad_(True).float() + base = list(pipe.vace.parameters()) + if not args.freeze_dit: + pipe.dit.requires_grad_(True).float() + base += list(pipe.dit.parameters()) + groups.append({"params": base, "lr": args.lr_base}) + if args.init_dit: + pipe.dit.load_state_dict(load_file(args.init_dit)) + if args.init_vace: + pipe.vace.load_state_dict(load_file(args.init_vace)) + if args.dit_lora_rank: + assert args.stage == "s2" and args.freeze_dit + groups.append({"params": lora.inject(pipe.dit, args.dit_lora_rank, args.dit_lora_rank), "lr": args.lora_lr}) + ema_params = [p for g in groups for p in g["params"]] # weights that end up in the checkpoints (not the repa projectors) + # SyncNet decode step: needed by the sync loss and by every loss on the decoded mouth / mouth tokens + aux = args.sync_weight > 0 or args.repa != "none" or args.pix_weight > 0 or args.lpips_weight > 0 + repa = None + if args.repa != "none": + repa = torch.nn.ModuleList([Repa(k, temporal=args.repa_temporal, cross=args.repa_cross).to(accel.device) for k in args.repa.split(",")]) + groups.append({"params": [p for r in repa for p in r.proj.parameters()], "lr": args.repa_lr}) + lpips_net = None + if args.lpips_weight > 0: + import lpips + torch.hub.set_dir(os.environ.get("VIDDUB_TORCH_HUB", str(PROJECT / "weights/torch_hub"))) + lpips_net = lpips.LPIPS(net="vgg", verbose=False).to(accel.device).eval().requires_grad_(False) + model = Trainable(pipe, adapter, SyncLoss().to(accel.device) if aux else None, repa, args.repa_layer, lpips_net) + opt = torch.optim.AdamW(groups, weight_decay=0.01) + + def lr_lambda(s): + if s < args.warmup: + return (s + 1) / args.warmup + if args.lr_schedule == "cosine": + return 0.1 + 0.45 * (1 + math.cos(math.pi * min(1.0, (s - args.warmup) / max(1, args.max_steps - args.warmup)))) + return 1.0 + sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda) + ema = [p.detach().float().clone() for p in ema_params] if args.ema > 0 else None + + sources = args.data.split(",") + weights = {k: float(v) for k, v in (kv.split("=") for kv in args.domain_weights.split(",") if kv)} + balanced = len(sources) > 1 or bool(weights) + cached = cached_clip_ids(sources) + train_rows = [r for r in load_rows("train", sources) if r["clip_id"] in cached][:args.max_train_clips] + clip_w = None + if args.sampling == "mix" and args.clip_weights: + cw = json.load(open(args.clip_weights)) + accel.print(f"clip weights {args.clip_weights}: keep {sum(r['clip_id'] in cw for r in train_rows)}/{len(train_rows)} train clips") + train_rows = [r for r in train_rows if r["clip_id"] in cw] + clip_w = [float(cw[r["clip_id"]]) for r in train_rows] + zero_latents = ZeroLatents(pipe) + zero_latents.empty_context() + dataset = CachedClipDataset(train_rows, args.num_frames, args.latents_name) + if balanced: + model, opt, sched = accel.prepare(model, opt, sched) + else: + loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=first_item, + num_workers=args.workers, persistent_workers=True, prefetch_factor=2) + model, opt, loader, sched = accel.prepare(model, opt, loader, sched) + accel.register_for_checkpointing(sched) + + step = 0 + state_dir = out / "state" + if state_dir.exists(): + accel.load_state(str(state_dir)) + step = json.load(open(out / "step.json"))["step"] + if ema is not None and (state_dir / "ema.pt").exists(): + ema = [t.to(accel.device) for t in torch.load(state_dir / "ema.pt")] + accel.print(f"resumed from step {step}") + if args.sampling == "epoch": + clips = {} + for r in train_rows: + clips[r["domain"]] = clips.get(r["domain"], 0) + 1 + stats = domain_stats(train_rows, {d: n * int(weights.get(d, 1)) for d, n in clips.items()}) + else: + stats = domain_stats(train_rows, weights) + if balanced: + if args.sampling == "mix": + sampler = MixSampler(train_rows, weights, clip_w or [1.0] * len(train_rows), accel.process_index, + accel.num_processes, step * args.grad_accum) + elif args.sampling == "epoch": + sampler = EpochSampler(train_rows, weights, accel.process_index, accel.num_processes, step * args.grad_accum) + accel.print(f"epoch sampling: {sampler.per * accel.num_processes} samples/epoch, " + f"{sampler.per / args.grad_accum:.0f} steps/epoch") + else: + sampler = DomainBalancedSampler(train_rows, weights, seed=1234 + accel.process_index + 1000003 * step) + loader = torch.utils.data.DataLoader(dataset, batch_size=1, sampler=sampler, collate_fn=first_item, + num_workers=args.workers, persistent_workers=True, prefetch_factor=2) + for d, s in stats.items(): + accel.print(f" {d:<11} clips {s['clips']:>7} groups {s['groups']:>6} share {s['share']:.3f} " + f"steps/epoch {s['clips'] / (s['share'] * accel.num_processes * args.grad_accum):.0f}") + + writer = SummaryWriter(str(out / "tb")) if accel.is_main_process else None + scorer, val_clips = None, None + if accel.is_main_process: + scorer = syncnet_scorer(accel.device) + val_rows = interleave_domains(load_rows("val", sources), args.val_clips) + val_clips = [load_clip(r, args.num_frames) for r in val_rows] + for r, c in zip(val_rows, val_clips): + c["lat"] = torch.load(PROJECT / "data/clips" / r["dir"] / args.latents_name) if r["clip_id"] in cached else None + c["domain"] = r["domain"] + json.dump(vars(args) | {"num_train_clips": len(train_rows), "world_size": accel.num_processes, "domains": stats, + "val_clip_ids": [r["clip_id"] for r in val_rows]}, open(out / "args.json", "w"), indent=1) + accel.print(f"train clips {len(train_rows)}, world {accel.num_processes}, trainable params " + f"{sum(p.numel() for g in groups for p in g['params']) / 1e6:.1f}M") + + raw = accel.unwrap_model(model) + t0, losses, sync_losses, pix_losses, repa_losses, lpips_losses, last_val = time.time(), [], [], [], [], [], (0 if args.skip_val0 else -1) + prev_audio = None + while step < args.max_steps: + for clip in loader: + if step % args.val_every == 0 and step != last_val: + if accel.is_main_process: + validate(args, accel, pipe, whisper, raw.adapter, zero_latents, val_clips, scorer, step, writer) + accel.wait_for_everyone() + last_val = step + + drop_ref = random.random() < args.p_drop_ref + x0, vace_context, context = build_inputs(clip["latents"], clip["mask"], args.num_frames, zero_latents, drop_ref=drop_ref, + context=None if random.random() < args.p_drop_text else clip["latents"]["context"]) + cross = aux and prev_audio is not None and random.random() < args.p_cross + audio = prev_audio if cross else clip["audio"] + prev_audio = clip["audio"] + hs = whisper([audio], accel.device, torch.bfloat16) + drop_audio = torch.tensor([not cross and random.random() < args.p_drop_audio]) + sync = None + if aux and not drop_audio.item(): + sync = {"box": mouth_box(clip["mask"], 0.1), "mel": mel_window(melspectrogram(audio), 1 / FPS).to(accel.device), + "weight": args.sync_weight, "sigma_max": args.sync_sigma_max, "cross": cross, "pix_weight": args.pix_weight, "lpips_weight": args.lpips_weight, + "repa_weight": args.repa_weight, "repa_detach": step < args.repa_warmup, "mel_full": melspectrogram(audio) if repa is not None else None} + with accel.accumulate(model), torch.autocast("cuda", dtype=torch.bfloat16): + loss, l_sync, l_pix, l_repa, l_lpips = model(x0, context, vace_context, hs, args.num_frames, 0 if drop_ref else 1, drop_audio, args.mask_lambda, sync) + accel.backward(loss) + if accel.sync_gradients: + norm = accel.clip_grad_norm_([p for g in groups for p in g["params"]], 1.0) + if not torch.isfinite(norm): # same on every rank (post all-reduce): skip the update everywhere + opt.zero_grad(set_to_none=True) + accel.print(f"step {step}: non-finite grad norm, update skipped", flush=True) + opt.step() + if accel.sync_gradients: + sched.step() + if ema is not None: + with torch.no_grad(): + torch._foreach_lerp_(ema, [p.detach().float() for p in ema_params], 1 - args.ema) + opt.zero_grad(set_to_none=True) + if l_sync is not None: + pipe.vae.model.clear_cache() # its feature cache keeps the decoder graph alive otherwise + sync_losses.append(l_sync.item()) + if l_pix is not None: + pix_losses.append(l_pix.item()) + if l_repa is not None: + repa_losses.append(l_repa.item()) + if l_lpips is not None: + lpips_losses.append(l_lpips.item()) + losses.append(loss.item()) + if not accel.sync_gradients: + continue + step += 1 + + if step % args.log_every == 0: + l = torch.tensor(np.mean(losses), device=accel.device) + l = accel.reduce(l, "mean").item() + ls = torch.tensor([sum(sync_losses), len(sync_losses)], device=accel.device, dtype=torch.float64) + ls = accel.reduce(ls, "sum") + ls = (ls[0] / ls[1].clamp(min=1)).item() + lp = accel.reduce(torch.tensor([sum(pix_losses), len(pix_losses)], device=accel.device, dtype=torch.float64), "sum") + lp = (lp[0] / lp[1].clamp(min=1)).item() + lr_ = accel.reduce(torch.tensor([sum(repa_losses), len(repa_losses)], device=accel.device, dtype=torch.float64), "sum") + lr_ = (lr_[0] / lr_[1].clamp(min=1)).item() + lq = accel.reduce(torch.tensor([sum(lpips_losses), len(lpips_losses)], device=accel.device, dtype=torch.float64), "sum") + lq = (lq[0] / lq[1].clamp(min=1)).item() + dt = (time.time() - t0) / args.log_every + accel.print(f"step {step} loss {l:.4f} sync {ls:.3f} pix {lp:.4f} repa {lr_:.4f} lpips {lq:.4f} lr {sched.get_last_lr()[0]:.2e} {dt:.2f}s/step " + f"mem {torch.cuda.max_memory_allocated() / 2**30:.1f}G", flush=True) + if writer: + writer.add_scalar("train/loss", l, step) + writer.add_scalar("train/sync", ls, step) + writer.add_scalar("train/pix", lp, step) + writer.add_scalar("train/repa", lr_, step) + writer.add_scalar("train/lpips", lq, step) + writer.add_scalar("train/sec_per_step", dt, step) + t0, losses, sync_losses, pix_losses, repa_losses, lpips_losses = time.time(), [], [], [], [], [] + + if step % args.save_every == 0 or step == args.max_steps: + accel.save_state(str(state_dir)) + if accel.is_main_process: + json.dump({"step": step}, open(out / "step.json", "w")) + save_weights(args, raw, out, step) + if ema is not None: + torch.save(ema, out / "state/ema.pt") + live = [p.detach().clone() for p in ema_params] + with torch.no_grad(): + torch._foreach_copy_(ema_params, ema) + (out.parent / f"{out.name}_ema").mkdir(exist_ok=True) + save_weights(args, raw, out.parent / f"{out.name}_ema", step) + with torch.no_grad(): + torch._foreach_copy_(ema_params, live) + del live + if step >= args.max_steps: + break + if accel.is_main_process: + validate(args, accel, pipe, whisper, raw.adapter, zero_latents, val_clips, scorer, step, writer) + accel.wait_for_everyone() + + +if __name__ == "__main__": + main()