scripts/inference/inference_streaming.py

#!/usr/bin/env python3
"""Streaming inference — per-chunk AR generation with decode-as-you-go.

Generates lip-synced video using the streaming pipeline: each AR chunk is
denoised, decoded, and composited before moving to the next. This enables
first-frame output before the full video is generated.

Supports three decoder modes:
  - StreamingTAEHV: temporal state across chunks (no boundary artifacts)
  - Batch TAEHV: independent per-chunk decode (faster, possible boundary seams)
  - Wan VAE: full VAE decode per chunk (highest quality, slowest)

The face detection + alignment + compositing pipeline always runs — streaming
has no raw-input path (use inference_segmentwise.py with --skip_preprocessing
for pre-aligned 512x512 inputs).

Usage:
    python scripts/inference/inference_streaming.py \
        --video_path /path/to/reference.mp4 \
        --output_path /path/to/output.mp4 \
        --ckpt_path /path/to/sf_trained_student.pth \
        --vae_path /path/to/Wan2.1_VAE.pth \
        --wav2vec_path /path/to/wav2vec2-base-960h \
        --mask_path /path/to/mask.png \
        --taehv_ckpt /path/to/taew2_1.pth \
        --text_embeds_path /path/to/text_emb.pt
"""

import argparse
import os

import cv2
import numpy as np
import torch
from PIL import Image

from _common import (
    TAEHVDecoderWrapper,
    load_vae, load_wav2vec, load_or_encode_text,
    resolve_audio, compute_generation_length,
    load_image_processor, preprocess_with_latentsync,
    build_condition, build_condition_from_precomputed,
    composite_with_latentsync_float,
    save_frames_as_video, mux_video_with_audio,
    encode_audio, frames_to_tensor, apply_spatial_mask, load_latentsync_mask,
    enumerate_samples,
)
# Use the 14B loader unconditionally — it dispatches on args.model_size
# (constructor_merge_lora flag + post-load PEFT merge gated to 14B).
# For 1.3B the LoRA-merge steps are skipped.
from _loader import load_diffusion_model


def parse_args():
    parser = argparse.ArgumentParser(
        description="Streaming inference with per-chunk decode.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # --- Model paths ---
    parser.add_argument("--ckpt_path", type=str, required=True,
                        help="SF-trained student checkpoint (.pth)")
    parser.add_argument("--vae_path", type=str, required=True,
                        help="Path to Wan2.1_VAE.pth")
    parser.add_argument("--wav2vec_path", type=str, required=True,
                        help="Path to wav2vec2-base-960h directory")
    parser.add_argument("--mask_path", type=str, required=True,
                        help="Path to LatentSync mask.png")
    parser.add_argument("--base_model_paths", type=str, default=None,
                        help="Comma-separated safetensor paths for base Wan 2.1 T2V "
                             "(not needed with the released self-contained checkpoint)")
    parser.add_argument("--model_size", type=str, default="14B",
                        choices=["1.3B", "14B"],
                        help="Student model size; 14B uses PEFT LoRA path.")
    parser.add_argument("--merge_lora_post_load", action="store_true", default=True,
                        help="14B only: merge PEFT LoRA into base after load_state_dict.")
    parser.add_argument("--no_merge_lora_post_load", action="store_false",
                        dest="merge_lora_post_load",
                        help="Disable post-load LoRA merge (keep PEFT layers active)")
    parser.add_argument("--omniavatar_ckpt_path", type=str, default=None,
                        help="OmniAvatar LoRA+audio checkpoint "
                             "(not needed with the released self-contained checkpoint)")
    parser.add_argument("--text_embeds_path", type=str, default=None,
                        help="Pre-computed T5 embeddings .pt file")
    parser.add_argument("--text_encoder_path", type=str, default=None,
                        help="T5 model path for runtime encoding")
    parser.add_argument("--prompt", type=str, default="a person talking",
                        help="Text prompt (encoded when --text_encoder_path is set)")

    # --- TAEHV ---
    parser.add_argument("--taehv_ckpt", type=str, default=None,
                        help="Path to TAEHV taew2_1.pth (required by the default "
                             "streaming_taehv / batch_taehv decoders)")

    # --- Streaming decoder mode ---
    parser.add_argument("--streaming_decoder", type=str, default="streaming_taehv",
                        choices=["streaming_taehv", "batch_taehv", "wan_vae"],
                        help="Decoder mode for streaming pipeline.")

    # --- Input/output ---
    parser.add_argument("--video_path", type=str, default=None,
                        help="Reference video path (any resolution)")
    parser.add_argument("--audio_path", type=str, default=None,
                        help="Separate audio source (extracted from video if not provided)")
    parser.add_argument("--output_path", type=str, default=None,
                        help="Output video path")
    parser.add_argument("--input_dir", type=str, default=None,
                        help="Directory of sample subdirs (each with sub_clip.mp4, audio.wav). "
                             "Mutually exclusive with --video_path.")
    parser.add_argument("--output_dir", type=str, default=None,
                        help="Output directory for batch mode")
    parser.add_argument("--skip_existing", action="store_true",
                        help="Skip samples whose output already exists (for resume)")

    # --- Generation params ---
    parser.add_argument("--t_list", nargs="+", type=float, default=[0.999, 0.769, 0.0],
                        help="Denoising schedule. Must match the checkpoint's distillation "
                             "schedule: the released 14B student is a 2-step t769 model "
                             "(0.999 -> 0.769 -> 0.0).")
    parser.add_argument("--chunk_size", type=int, default=3,
                        help="Number of latent frames per AR chunk")
    parser.add_argument("--num_latent_frames", type=int, default=None,
                        help="Override generation length (must be multiple of chunk_size)")
    parser.add_argument("--min_latent_frames", type=int, default=0,
                        help="Floor on num_latent; shorter audio is padded via zero-audio "
                             "+ ping-pong video. 0 disables.")
    parser.add_argument("--context_noise", type=float, default=0.0,
                        help="Noise added to context frames during AR generation")
    parser.add_argument("--seed", type=int, default=42,
                        help="Random seed")
    parser.add_argument("--fps", type=float, default=25.0,
                        help="Output video FPS")
    parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16", "fp32"],
                        help="Model dtype")
    parser.add_argument("--device", type=str, default="cuda",
                        help="Device for inference")

    # --- Attention ---
    parser.add_argument("--local_attn_size", type=int, default=7,
                        help="Rolling local attention window in latent frames. Default 7 "
                             "(the trained window: 1 sink + 6 rolling) keeps VRAM constant "
                             "for any clip length. -1 = full attention over the whole clip "
                             "(VRAM grows with clip length).")
    parser.add_argument("--sink_size", type=int, default=1,
                        help="Number of initial latent frames always kept in the attention "
                             "window (default 1, matching training)")
    parser.add_argument("--use_dynamic_rope", action="store_true", default=True,
                        help="Window-local dynamic RoPE (default: on, matching training)")
    parser.add_argument("--no_dynamic_rope", action="store_false", dest="use_dynamic_rope",
                        help="Disable window-local dynamic RoPE (absolute positions; "
                             "pair with --local_attn_size -1)")

    # --- Preprocessing (face detection + alignment + compositing) ---
    parser.add_argument("--skip_preprocessing", action="store_true",
                        help="NOT SUPPORTED by the streaming pipeline (it composites "
                             "per chunk); use inference_segmentwise.py for pre-aligned "
                             "512x512 inputs. Passing this flag raises an error.")
    parser.add_argument("--face_cache_dir", type=str, default=None,
                        help="Optional directory for face-detection caches; speeds up "
                             "repeated runs over the same videos. Unset = no caching.")
    parser.add_argument("--composite_full_face", action="store_true",
                        help="Composite the entire generated 512x512 face back into the "
                             "original frame. Default: blend only the mouth region of the "
                             "generated face; the rest stays from the input video.")

    # --- Streamwise encoding (truly interleaved encode/denoise/decode) ---
    parser.add_argument("--streamwise_encode", action="store_true", default=True,
                        help="Encode the source video chunk-by-chunk inside the AR "
                             "loop (encoder feat_cache preserved across chunks), so "
                             "GPU memory stays constant for any clip length. Output "
                             "is bit-identical to the upfront full-clip encode. "
                             "Default: on.")
    parser.add_argument("--no_streamwise_encode", action="store_false",
                        dest="streamwise_encode",
                        help="Encode the whole reference video upfront instead "
                             "(adds roughly 2.4 GB GPU memory per minute of input).")

    # --- Deferred compositing (move lip-blend + affine warp out of AR loop) ---
    parser.add_argument("--defer_composite", action="store_true",
                        help="Skip per-block compositing inside the AR loop; "
                             "concat all decoded chunks and run "
                             "composite_with_latentsync_float once after the "
                             "loop ends. Improves throughput (no per-block "
                             ".cpu() sync) at the cost of first-frame latency.")

    # --- torch.compile ---
    parser.add_argument("--compile", action="store_true",
                        help="Wrap diffusion model + Wan VAE encoder/decoder + "
                             "TAEHV (when present) with torch.compile.")

    return parser.parse_args()


def validate_args(args):
    if args.input_dir is not None and args.video_path is not None:
        raise ValueError("--input_dir and --video_path are mutually exclusive")
    if args.input_dir is None and args.video_path is None:
        raise ValueError("Must provide either --input_dir or --video_path")
    if args.input_dir is not None and args.output_dir is None:
        raise ValueError("--input_dir requires --output_dir")
    if args.input_dir is None and args.output_path is None:
        raise ValueError("--video_path mode requires --output_path")
    if args.skip_preprocessing:
        raise ValueError(
            "--skip_preprocessing is not supported by the streaming pipeline "
            "(it composites each decoded chunk back into the original frames). "
            "Use inference_segmentwise.py for pre-aligned 512x512 inputs."
        )
    if args.text_embeds_path is None and args.text_encoder_path is None:
        raise ValueError(
            "Text conditioning is required: pass --text_embeds_path "
            "(precomputed T5 embeddings) or --text_encoder_path "
            "(encodes --prompt at runtime)."
        )
    if args.streaming_decoder in ("streaming_taehv", "batch_taehv") and not args.taehv_ckpt:
        raise ValueError(f"--streaming_decoder {args.streaming_decoder} requires --taehv_ckpt")


def build_condition_streamwise(vae, wav2vec_model, wav2vec_extractor,
                                video_frames_np, audio_path, text_embeds,
                                mask_path, num_video_frames, num_latent_frames,
                                device, dtype):
    """Build a *minimal* condition dict for streamwise AR inference.

    Encodes only audio (full upfront) and the very first ref_latent (1 frame).
    Returns the condition with ref_sequence/masked_latents = None plus the
    pixel-space video tensors that the AR loop will encode incrementally.
    """
    # ============================================================
    # STAGE 2: Wav2Vec2 audio encode (full audio at once)
    # ============================================================
    print("Encoding audio (full) ...")
    audio_emb = encode_audio(
        wav2vec_model, wav2vec_extractor, audio_path, num_video_frames, device
    )
    audio_emb = audio_emb.to(dtype=dtype)

    # ============================================================
    # STAGE 3a (streamwise stub): pixel-space tensors only, no encode
    # ============================================================
    # The reference video stays in pixel space here. STAGE 3b inside the
    # AR loop encodes it chunk-by-chunk via streaming_encode_chunk.
    H, W = 512, 512
    video_tensor = frames_to_tensor(video_frames_np)  # [1, 3, T, H, W] in [-1, 1]

    mask_img = Image.open(mask_path).convert("L")
    mask_np = np.array(mask_img).astype(np.float32) / 255.0
    if mask_np.shape[0] != H or mask_np.shape[1] != W:
        mask_np = cv2.resize(mask_np, (W, H), interpolation=cv2.INTER_LINEAR)
    mask_pixel_binary = (mask_np > 0.5).astype(np.float32)
    masked_video_tensor = apply_spatial_mask(
        video_tensor, mask_pixel_binary, mask_all_frames=True
    )

    # No VAE encoding here. The AR loop in run_streaming_pipeline will encode
    # both the unmasked (ref) and masked streams chunk-by-chunk, in lockstep
    # with denoise + decode.
    H_lat = H // 8
    W_lat = W // 8
    latent_mask = load_latentsync_mask(mask_path, H_lat, W_lat).to(device=device, dtype=dtype)

    condition = {
        "text_embeds": text_embeds,
        "audio_emb": audio_emb,
        "ref_latent": None,       # set on block 0 from growing_ref_seq[..., :1]
        "mask": latent_mask.to(device=device),
        "masked_video": None,   # built incrementally in AR loop
        "ref_sequence": None,     # built incrementally in AR loop
    }
    return condition, video_tensor, masked_video_tensor


@torch.no_grad()
def run_streaming_pipeline(
    model, decoder_vae, vae, condition, num_latent_frames, num_video_frames,
    args, latentsync_metadata, image_processor, audio_path, output_path,
    device, dtype,
    video_tensor=None, masked_video_tensor=None,
):
    """Run the streaming pipeline: per-chunk denoise → decode → composite.

    Stage map (stage labels used by the LatentSync compositing path):
        STAGE 1  : Face detect + 512x512 alignment    (handled before this
                   function, in main() via preprocess_with_latentsync)
        STAGE 2  : Wav2Vec2 audio encode               (handled in build_condition[_streamwise])
        STAGE 3a : Reference VAE encode (full or first-frame only)
        STAGE 3b : (streamwise only) per-block VAE encode of unmasked + masked
        STAGE 4  : Per-block 2-step DDIM denoise (CausalOmniAvatarWan)
        STAGE 5  : Per-block VAE decode (Wan VAE / TAEHV / StreamingTAEHV)
        STAGE 6  : Per-block compositing (paste lip region back into full-res frame)
        STAGE 7  : Per-block KV cache update (model forward with denoised x0 stored as cache)
        STAGE 8  : (streaming_taehv only) flush remaining buffered frames
        STAGE 9  : Save MP4 + ffmpeg audio mux

    Returns:
        composited_np: [N, H, W, 3] uint8 numpy array of composited frames.
    """
    # ============================================================
    # STAGE 5-prep: Decoder selection
    # ============================================================
    _use_streaming_dec = False
    if args.streaming_decoder == "streaming_taehv":
        try:
            from lipforcing.methods.reward.taehv import StreamingTAEHV
            if hasattr(decoder_vae, 'taehv') and decoder_vae.taehv is not None:
                streaming_dec = StreamingTAEHV(decoder_vae.taehv)
                _use_streaming_dec = True
                print("  Using StreamingTAEHV decoder (temporal state across chunks)")
        except Exception:
            pass
    if not _use_streaming_dec:
        if args.streaming_decoder == "wan_vae":
            print("  Using Wan VAE decoder per chunk")
        else:
            print("  Using batch TAEHV decoder per chunk")

    # ============================================================
    # VRAM management: offload encoder to CPU when only the decoder is needed
    # ============================================================
    # In streamwise_encode mode the encoder is needed inside the AR loop, so
    # keep it on GPU even when the decoder is a different module (e.g. TAEHV).
    if (decoder_vae is not vae and hasattr(vae, 'parameters')
            and not args.streamwise_encode):
        vae.to("cpu")
    torch.cuda.empty_cache()

    # ============================================================
    # STAGE 5-prep: Wan VAE streaming-decode cache reset
    # ============================================================
    # For wan_vae streaming, reset feat_cache once before AR loop. Subsequent
    # streaming_decode_chunk() calls preserve cache across chunks so output
    # at chunk boundaries is bit-identical to a single full-length decode.
    _use_wan_streaming = (args.streaming_decoder == "wan_vae"
                          and hasattr(decoder_vae, "streaming_decode_chunk"))
    if _use_wan_streaming:
        print("  Wan VAE: streaming-decode mode (cache continuity across chunks)")
        decoder_vae.reset_decode_cache()

    # ============================================================
    # STAGE 3b-prep: Streamwise encode setup (only if --streamwise_encode)
    # ============================================================
    _streamwise = args.streamwise_encode and video_tensor is not None
    if _streamwise:
        print("  Wan VAE: streamwise-encode mode (encoder in AR loop)")
        # Two independent encoder feat_cache streams: one for unmasked
        # (ref_sequence) and one for masked (masked_latents). We share a
        # single VAE instance and swap cache state between calls.
        vae.reset_encode_cache()
        unmasked_state = vae.save_encode_cache_state()  # both empty
        masked_state = vae.save_encode_cache_state()
        growing_ref_seq = None     # [1, 16, T_so_far, H_lat, W_lat]
        growing_masked = None
        original_vae_dtype = next(vae.parameters()).dtype
        vae.to(dtype=torch.bfloat16)

        def _frame_chunks_for_block(block_idx):
            """Return list of (start, end) frame indices for an AR block.
            Block 0: 1 + 4 + 4 = 9 frames (3 latents).
            Block i>=1: 4 + 4 + 4 = 12 frames (3 latents).
            """
            if block_idx == 0:
                return [(0, 1), (1, 5), (5, 9)]
            base = 9 + 12 * (block_idx - 1)
            return [(base, base + 4), (base + 4, base + 8), (base + 8, base + 12)]

        def _stream_encode_block(block_idx, source_video_tensor, prev_state):
            """Encode the next AR block's worth of frames into 3 latents."""
            vae.load_encode_cache_state(prev_state)
            chunks = []
            for s, e in _frame_chunks_for_block(block_idx):
                chunk = source_video_tensor[0, :, s:e].to(
                    dtype=torch.bfloat16, device=device)
                chunks.append(vae.streaming_encode_chunk(chunk, device=device))
            new_state = vae.save_encode_cache_state()
            # chunks are [1, 16, 1, H, W]; concat along time
            new_lats = torch.cat([c.squeeze(0) for c in chunks], dim=1).unsqueeze(0)
            return new_lats.to(dtype=dtype), new_state

    # --- Prepare model ---
    model.total_num_frames = num_latent_frames
    model.clear_caches()
    B, C = 1, 16
    if condition.get("ref_latent") is not None:
        H_lat, W_lat = condition["ref_latent"].shape[3], condition["ref_latent"].shape[4]
    else:
        # streamwise mode: derive from pixel-space video tensor (8x spatial compression)
        H_lat = video_tensor.shape[-2] // 8
        W_lat = video_tensor.shape[-1] // 8
    t_list_t = torch.tensor(args.t_list, device=device, dtype=torch.float64)

    # Pre-generate all noise at once (must match non-streaming pipeline)
    torch.manual_seed(args.seed)
    all_noise = torch.randn(B, C, num_latent_frames, H_lat, W_lat, device=device, dtype=dtype)

    num_blocks = num_latent_frames // args.chunk_size
    all_composited_frames = []
    # When --defer_composite is on, we skip per-block compositing and stash
    # the raw decoded chunk_float tensors here, then composite once after
    # the AR loop ends.
    all_decoded_chunks_cpu = []
    video_frame_offset = 0

    # ============================================================
    # STAGE 3b/4/5/6/7: AR streaming loop (repeats num_blocks times)
    # ============================================================
    for block_idx in range(num_blocks):
        cur_start_frame = block_idx * args.chunk_size

        # ----------------------------------------------------------
        # STAGE 3b: per-block VAE encode (streamwise mode only)
        # ----------------------------------------------------------
        # Block 0 ingests 9 video frames (1+4+4) -> 3 latents.
        # Subsequent blocks ingest 12 frames (4+4+4) -> 3 latents.
        # Two encoder feat_cache streams are swapped (unmasked vs masked) so
        # each stream maintains its own continuous temporal context across
        # all blocks.
        if _streamwise:
            new_unmasked, unmasked_state = _stream_encode_block(
                block_idx, video_tensor, unmasked_state)
            new_masked, masked_state = _stream_encode_block(
                block_idx, masked_video_tensor, masked_state)

            growing_ref_seq = (new_unmasked if growing_ref_seq is None
                               else torch.cat([growing_ref_seq, new_unmasked], dim=2))
            growing_masked = (new_masked if growing_masked is None
                              else torch.cat([growing_masked, new_masked], dim=2))

            condition["ref_sequence"] = growing_ref_seq
            condition["masked_video"] = growing_masked
            condition["ref_latent"] = growing_ref_seq[:, :, :1].contiguous()

        # ----------------------------------------------------------
        # STAGE 4: 2-step DDIM denoise on this block's 3 latents
        # ----------------------------------------------------------
        # t_list = [0.999, 0.833, 0]; len(t_list)-1 = 2 model forwards per
        # block. Self-attention is causal sliding-window (sink=1, window=7
        # AR chunks); cross-attention attends to audio_emb + ref_sequence.
        noisy_input = all_noise[:, :, cur_start_frame:cur_start_frame + args.chunk_size]
        for step_idx in range(len(t_list_t) - 1):
            t_cur = t_list_t[step_idx]
            t_next = t_list_t[step_idx + 1]
            x0_pred = model(
                noisy_input, t_cur.expand(B),
                condition=condition,
                cur_start_frame=cur_start_frame,
                store_kv=False, is_ar=True,
                fwd_pred_type="x0", use_gradient_checkpointing=False,
            )
            if t_next > 0:
                eps = torch.randn_like(x0_pred)
                noisy_input = model.noise_scheduler.forward_process(
                    x0_pred, eps, t_next.expand(B))
            else:
                noisy_input = x0_pred

        # ----------------------------------------------------------
        # STAGE 5: per-block VAE decode (3 latents -> ~12 video frames)
        # ----------------------------------------------------------
        # Three decoder modes:
        #   - StreamingTAEHV: per-latent decode with MemBlock temporal state
        #     across chunks; first chunk emits fewer frames (buffering).
        #   - Wan VAE streaming: per-latent decode_chunk with feat_cache
        #     persistence; bit-identical to single full-length decode.
        #   - batch TAEHV / Wan VAE batch: 3 latents at once, no continuity.
        if _use_streaming_dec:
            chunk_latent = x0_pred[0].to(device, dtype=torch.float16)
            chunk_latent_ntchw = chunk_latent.permute(1, 0, 2, 3).unsqueeze(0)
            chunk_frames = []
            for t in range(chunk_latent_ntchw.shape[1]):
                latent_t = chunk_latent_ntchw[:, t:t+1]
                frame = streaming_dec.decode(latent_t)
                while frame is not None:
                    chunk_frames.append(frame)
                    frame = streaming_dec.decode()
            if chunk_frames:
                chunk_float = torch.cat(chunk_frames, dim=1).squeeze(0)
            else:
                chunk_float = None
        elif _use_wan_streaming:
            # Stream one latent at a time so feat_cache state advances
            # exactly as in the per-latent inner loop of VideoVAE_.decode.
            vae_dtype = next(decoder_vae.parameters()).dtype
            chunk_latent = x0_pred[0].to(vae_dtype)  # [16, 3, h_lat, w_lat]
            video_chunks = []
            for t in range(chunk_latent.shape[1]):
                latent_t = chunk_latent[:, t:t+1]  # [16, 1, h_lat, w_lat]
                v = decoder_vae.streaming_decode_chunk(latent_t, device=device)
                # v: [1, 3, t_video, H, W] in [-1, 1]; t_video = 1 on the very
                # first call, 4 thereafter (Wan VAE 4x temporal upsampling).
                video_chunks.append(v)
            chunk_decoded = torch.cat(video_chunks, dim=2)  # [1, 3, T, H, W]
            chunk_float = chunk_decoded[0].permute(1, 0, 2, 3)
            chunk_float = ((chunk_float + 1) / 2).clamp(0, 1)
        else:
            chunk_latent = x0_pred[0].to(torch.float32)
            chunk_decoded = decoder_vae.decode([chunk_latent], device=device)
            chunk_decoded = chunk_decoded.clamp(-1, 1)
            chunk_float = chunk_decoded[0].permute(1, 0, 2, 3)
            chunk_float = ((chunk_float + 1) / 2).clamp(0, 1)

        # ----------------------------------------------------------
        # STAGE 6: per-block compositing (CPU)
        # ----------------------------------------------------------
        # Paste the generated 512x512 lip region back into the full-res
        # frame using the LatentSync affine matrices captured during
        # Stage 1. CPU-bound; runs after each chunk_float arrives.
        # When --defer_composite is on, skip this and stash the raw
        # decoded chunk for one batch composite after the AR loop.
        if chunk_float is not None:
            if args.defer_composite:
                all_decoded_chunks_cpu.append(chunk_float.cpu())
                video_frame_offset += all_decoded_chunks_cpu[-1].shape[0]
            else:
                composited = composite_with_latentsync_float(
                    chunk_float.cpu(), latentsync_metadata, image_processor,
                    use_mouth_only_compositing=not args.composite_full_face,
                    frame_offset=video_frame_offset,
                )
                all_composited_frames.append(composited)
                video_frame_offset += composited.shape[0]

        # ----------------------------------------------------------
        # STAGE 7: KV cache update (extra model forward per block)
        # ----------------------------------------------------------
        # Re-run the model with this block's denoised x0_pred (or a
        # noised version when context_noise > 0) and store_kv=True so
        # subsequent blocks have valid sliding-window self-attention
        # context. This is the AR carry that makes the next block's
        # generation conditioned on the past.
        cache_input = x0_pred
        t_cache = torch.full((B,), args.context_noise, device=device, dtype=torch.float64)
        if args.context_noise > 0:
            cache_eps = torch.randn_like(x0_pred)
            cache_input = model.noise_scheduler.forward_process(
                x0_pred, cache_eps,
                torch.tensor(args.context_noise, device=device, dtype=torch.float64).expand(B))
        model(cache_input, t_cache, condition=condition,
              cur_start_frame=cur_start_frame, store_kv=True, is_ar=True,
              fwd_pred_type="x0", use_gradient_checkpointing=False)

        if (block_idx + 1) % 5 == 0 or block_idx == num_blocks - 1:
            print(f"  Streaming block {block_idx + 1}/{num_blocks} done")

    model.clear_caches()

    if _streamwise:
        vae.to(dtype=original_vae_dtype)

    # ============================================================
    # STAGE 8: Flush remaining buffered frames (streaming_taehv only)
    # ============================================================
    # StreamingTAEHV needs future temporal context to emit frames, so
    # the very last latents stay buffered until we explicitly flush at
    # the end of the AR loop.
    if _use_streaming_dec:
        flush_frames = streaming_dec.flush_decoder()
        if flush_frames:
            flush_float = torch.cat(flush_frames, dim=1).squeeze(0)
            flush_cpu = flush_float.cpu()
            if args.defer_composite:
                all_decoded_chunks_cpu.append(flush_cpu)
                video_frame_offset += flush_cpu.shape[0]
            else:
                composited = composite_with_latentsync_float(
                    flush_cpu, latentsync_metadata, image_processor,
                    use_mouth_only_compositing=not args.composite_full_face,
                    frame_offset=video_frame_offset,
                )
                all_composited_frames.append(composited)
                video_frame_offset += composited.shape[0]

    # ============================================================
    # STAGE 6 (deferred): one-shot compositing over all decoded frames
    # ============================================================
    # Only runs when --defer_composite is on. The lip blend + affine
    # warp are per-frame ops with no temporal coupling, so doing them
    # in one batch is identical to doing them per chunk -- but lets
    # the GPU run the whole AR loop without breaking pipelining on
    # per-block .cpu() syncs.
    if args.defer_composite and all_decoded_chunks_cpu:
        all_decoded = torch.cat(all_decoded_chunks_cpu, dim=0)
        composited_np = composite_with_latentsync_float(
            all_decoded, latentsync_metadata, image_processor,
            use_mouth_only_compositing=not args.composite_full_face,
            frame_offset=0,
        )
    else:
        composited_np = np.concatenate(all_composited_frames, axis=0)

    # ============================================================
    # STAGE 9: Save MP4 + ffmpeg audio mux (CPU)
    # ============================================================
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    save_frames_as_video(composited_np, output_path, fps=args.fps)
    video_duration = composited_np.shape[0] / args.fps
    tmp_composited = output_path + ".tmp.mp4"
    os.rename(output_path, tmp_composited)
    mux_video_with_audio(tmp_composited, audio_path, output_path, duration_s=video_duration)
    if os.path.exists(tmp_composited):
        os.remove(tmp_composited)

    return composited_np


def main():
    args = parse_args()
    validate_args(args)

    # Activate @conditional_compile decorators in network_causal.py BEFORE
    # the model class is imported (which happens later via load_diffusion_model).
    if args.compile:
        os.environ["LIPFORCING_COMPILE"] = "true"

    dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
    dtype = dtype_map[args.dtype]
    device = torch.device(args.device)

    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    # --- Load models ---
    print("Loading diffusion model ...")
    model = load_diffusion_model(args, device, dtype)

    print("Loading VAE ...")
    vae = load_vae(args.vae_path, device)

    # Decoder selection based on streaming_decoder mode
    if args.streaming_decoder in ("streaming_taehv", "batch_taehv"):
        if not args.taehv_ckpt:
            raise ValueError(f"--streaming_decoder {args.streaming_decoder} requires --taehv_ckpt")
        print(f"Loading TAEHV decoder from {args.taehv_ckpt} ...")
        decoder_vae = TAEHVDecoderWrapper(args.taehv_ckpt, device)
    else:
        decoder_vae = vae

    encoder_vae = vae

    # Eagerly load Wav2Vec + text
    print("Loading Wav2Vec2 (eager) ...")
    wav2vec_model, wav2vec_extractor = load_wav2vec(args.wav2vec_path, device)
    # OmniAvatar Wav2VecModel requires seq_len + output_hidden_states.
    _dummy_audio = np.zeros(16000, dtype=np.float32)
    _dummy_input = wav2vec_extractor(_dummy_audio, return_tensors="pt", sampling_rate=16000)
    with torch.no_grad():
        wav2vec_model(
            _dummy_input.input_values.to(device),
            seq_len=25, output_hidden_states=True,
        )
    print("Wav2Vec2 warmed up.")

    print("Loading text embeddings ...")
    text_embeds = load_or_encode_text(args, device, dtype)

    # LatentSync ImageProcessor (face detection + alignment; always on —
    # the streaming pipeline composites every decoded chunk).
    image_processor = load_image_processor(args.mask_path, device)

    # ===================================================================
    # Optional torch.compile wrapping (compile time absorbed by warmup)
    # ===================================================================
    if args.compile:
        # Compile activated via @conditional_compile decorators on hot
        # functions (see lipforcing/networks/OmniAvatar/inference_utils.py).
        # The env var that activates them was set at top of main() before
        # the model was imported.
        print("[--compile] Hot functions decorated with @conditional_compile.")

        # Compile Wan VAE encode/decode paths.
        # TAEHV is skipped — its internals do enumerate(self.decoder)
        # which breaks when Sequential is wrapped in OptimizedModule.
        _compile_kw = dict(mode=None, backend="inductor", dynamic=None)

        # Wan VAE decoder compile (skip if decoder is TAEHV)
        if not isinstance(decoder_vae, TAEHVDecoderWrapper):
            if hasattr(decoder_vae, 'model') and hasattr(decoder_vae.model, 'decoder'):
                decoder_vae.model.decoder = torch.compile(
                    decoder_vae.model.decoder, **_compile_kw)
                print("[--compile] Wan VAE decoder compiled.")

        # Wan VAE encoder compile
        if hasattr(encoder_vae, 'model') and hasattr(encoder_vae.model, 'encoder'):
            encoder_vae.model.encoder = torch.compile(
                encoder_vae.model.encoder, **_compile_kw)
            print("[--compile] Wan VAE encoder compiled.")

    # --- Loop over samples ---
    samples = list(enumerate_samples(args))
    succeeded, failed, skipped = [], [], []

    for sample_idx, (name, video_path, audio_path_sample, precomputed_dir) in enumerate(samples):
        print(f"\n{'='*60}")
        print(f"[{sample_idx+1}/{len(samples)}] {name}")
        print(f"{'='*60}")

        if args.input_dir is not None:
            output_path = os.path.join(args.output_dir, f"{name}.mp4")
        else:
            output_path = args.output_path

        if args.skip_existing and os.path.isfile(output_path):
            print(f"  [Skip] {output_path}")
            skipped.append(name)
            continue

        tmp_audio = None
        try:
            audio_path, tmp_audio = resolve_audio(
                audio_path=audio_path_sample, video_path=video_path,
            )

            num_latent_frames, num_video_frames = compute_generation_length(
                audio_path, args.num_latent_frames, args.chunk_size, args.fps,
                min_latent_frames=args.min_latent_frames,
            )

            # ============================================================
            # STAGE 1: Face detect + 512x512 affine alignment (CPU+GPU)
            # ============================================================
            # InsightFace (buffalo_l) bounding box detection followed by
            # LatentSync's affine_transform crop. Returns aligned 512x512
            # face crops + per-frame affine matrices for paste-back.
            print("Running LatentSync face detection ...")
            latentsync_metadata = preprocess_with_latentsync(
                video_path, image_processor, args.face_cache_dir,
                num_frames=num_video_frames,
            )
            if latentsync_metadata is None:
                print("  [FAIL] LatentSync preprocessing failed")
                failed.append(name)
                continue

            # Build conditioning
            if precomputed_dir is not None:
                condition = build_condition_from_precomputed(
                    precomputed_dir, args.mask_path,
                    num_latent_frames, device, dtype,
                )
                video_tensor = masked_video_tensor = None
            else:
                aligned_faces = latentsync_metadata["aligned_faces"]
                ref_frames_np = np.stack([
                    f.permute(1, 2, 0).numpy() if isinstance(f, torch.Tensor) else f
                    for f in aligned_faces[:num_video_frames]
                ], axis=0)

                # ============================================================
                # STAGE 2 + 3a: Audio encode + reference VAE encode (build condition)
                # ============================================================
                # Streamwise mode: only audio_emb is computed here; the
                # reference video is encoded inside the AR loop (STAGE 3b).
                # Default mode: full reference video VAE-encoded upfront.
                if args.streamwise_encode:
                    print("Building conditioning (streamwise) ...")
                    condition, video_tensor, masked_video_tensor = (
                        build_condition_streamwise(
                            encoder_vae, wav2vec_model, wav2vec_extractor,
                            ref_frames_np, audio_path, text_embeds, args.mask_path,
                            num_video_frames, num_latent_frames, device, dtype,
                        )
                    )
                else:
                    print("Building conditioning ...")
                    condition = build_condition(
                        encoder_vae, wav2vec_model, wav2vec_extractor, ref_frames_np,
                        audio_path, text_embeds, args.mask_path,
                        num_video_frames, num_latent_frames, device, dtype,
                    )
                    video_tensor = masked_video_tensor = None

            # Run streaming pipeline
            print(f"Running streaming pipeline ({args.streaming_decoder}) ...")
            run_streaming_pipeline(
                model, decoder_vae, vae, condition,
                num_latent_frames, num_video_frames,
                args, latentsync_metadata, image_processor,
                audio_path, output_path, device, dtype,
                video_tensor=video_tensor,
                masked_video_tensor=masked_video_tensor,
            )

            succeeded.append(name)
            print(f"  Done: {output_path}")

        except Exception as e:
            print(f"  [ERROR] {name}: {e}")
            import traceback
            traceback.print_exc()
            failed.append(name)
        finally:
            if tmp_audio is not None and os.path.exists(tmp_audio):
                os.remove(tmp_audio)
            torch.cuda.empty_cache()

    # --- Summary ---
    print(f"\n{'='*60}")
    print(f"Summary: {len(succeeded)} succeeded, {len(failed)} failed, {len(skipped)} skipped")
    if failed:
        print(f"  Failed: {failed}")



if __name__ == "__main__":
    main()