scripts/inference/inference_segmentwise.py

#!/usr/bin/env python3
"""Segment-wise inference — block-wise AR generation, decode after the full rollout.

Generates lip-synced video from a reference video and audio using the
CausalOmniAvatarWan student model (14B by default; ``--model_size 1.3B`` also
supported) trained via Self-Forcing. The whole clip is denoised chunk-by-chunk
first, then decoded and composited in one pass — the maximum-quality path
(use ``inference_streaming.py`` for decode-as-you-go / low-latency output).

Decoding uses the full Wan VAE by default; pass ``--taehv_ckpt`` to swap in
the TAEHV tiny decoder for throughput (``--taehv_streaming`` /
``--taehv_encode`` further extend that).

Face detection + alignment + compositing (the LatentSync preprocessing
pipeline) runs by default so arbitrary talking-head videos work as input.
Pass ``--skip_preprocessing`` only when inputs are already 512x512 aligned
face crops.

Usage:
    python scripts/inference/inference_segmentwise.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 \
        --text_embeds_path /path/to/text_emb.pt
"""

import argparse
import os

import numpy as np
import torch

from _common import (
    TAEHVDecoderWrapper, StreamingTAEHVDecoderWrapper,
    load_vae, load_wav2vec, load_or_encode_text,
    resolve_audio, compute_generation_length,
    load_and_adjust_video,
    load_image_processor, preprocess_with_latentsync,
    build_condition, build_condition_from_precomputed,
    composite_with_latentsync_float,
    save_frames_as_video, mux_video_with_audio,
    enumerate_samples,
)
from _loader import load_diffusion_model  # (shared; see scripts/inference/_loader.py)


# ===========================================================================
# CLI argument parsing
# ===========================================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description="Segment-wise causal OmniAvatar inference (block-wise AR with audio conditioning)"
    )

    # --- Single-sample mode ---
    parser.add_argument("--video_path", type=str, default=None,
                        help="Reference video path (any resolution; must be 512x512 "
                             "only with --skip_preprocessing)")
    parser.add_argument("--output_path", type=str, default=None,
                        help="Output video path")
    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("--taehv_ckpt", type=str, default=None,
                        help="Optional path to TAEHV taew2_1.pth. If set, uses the TAEHV "
                             "tiny decoder for latent->pixel decoding (full Wan VAE is still "
                             "used for encoding driving video unless --taehv_encode is set).")
    parser.add_argument("--taehv_encode", action="store_true",
                        help="Also use TAEHV for encoding the driving video (requires --taehv_ckpt). "
                             "Default: full Wan VAE encoder.")
    parser.add_argument("--taehv_streaming", action="store_true",
                        help="Use StreamingTAEHV for decoding (feeds latents one at a time). "
                             "Requires --taehv_ckpt.")
    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")

    # --- Optional model paths ---
    parser.add_argument("--base_model_paths", type=str, default=None,
                        help="Comma-separated safetensor paths for base Wan 2.1 T2V (1.3B or 14B)")
    parser.add_argument("--omniavatar_ckpt_path", type=str, default=None,
                        help="OmniAvatar LoRA+audio checkpoint")
    parser.add_argument("--audio_path", type=str, default=None,
                        help="Separate audio source (extracted from video if not provided)")

    # --- Generation parameters ---
    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; if audio is shorter, pad via zero-audio + ping-pong "
                             "video. 0 disables. 21 corresponds to the 81-frame (21 latent) generation length.")
    parser.add_argument("--prompt", type=str, default="a person talking",
                        help="Text prompt")
    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("--precomputed_dir", type=str, default=None,
                        help="Directory with precomputed .pt files (vae_latents_mask_all.pt, "
                             "ref_latents.pt, audio_emb_omniavatar.pt, text_emb.pt). "
                             "Bypasses VAE/Wav2Vec encoding — uses exact training-style tensors.")

    # --- Batch inference ---
    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. Training-style sample "
                             "dirs contain pre-aligned 512x512 crops — combine with "
                             "--skip_preprocessing (face detection fails on tight crops).")
    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)")

    # --- Preprocessing (face detection + alignment + compositing) ---
    parser.add_argument("--skip_preprocessing", action="store_true",
                        help="Skip the face detection + 512x512 alignment + compositing "
                             "pipeline. Requires inputs that are ALREADY 512x512 aligned "
                             "face crops; the output is the raw generated video (no "
                             "paste-back into the original frames).")
    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.")
    parser.add_argument("--save_aligned", action="store_true",
                        help="Additionally save the raw generated 512x512 aligned face "
                             "video as <output>_aligned.mp4 (before compositing).")

    parser.add_argument("--t_list", type=float, nargs="+",
                        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("--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)")
    parser.add_argument("--model_size", type=str, default="14B",
                        choices=["1.3B", "14B"],
                        help="Student model size. 14B is the default for SF LoRA runs.")
    parser.add_argument("--merge_lora_post_load", action="store_true", default=True,
                        help="After loading the SF trainable LoRA values, merge them into "
                             "base for inference speed. The model is constructed with "
                             "merge_lora=False (to expose lora_A/lora_B keys for the "
                             "trainable-filtered SF state_dict), then merged in-place "
                             "after load_state_dict.  Set --no_merge_lora_post_load to keep "
                             "PEFT layers active (slower forward, useful for debugging).")
    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("--chunk_size", type=int, default=3,
                        help="Number of latent frames per AR chunk")
    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("--device", type=str, default="cuda",
                        help="Device for inference")
    parser.add_argument("--dtype", type=str, default="bf16",
                        choices=["bf16", "fp16", "fp32"],
                        help="Model dtype")
    parser.add_argument("--fps", type=int, default=25,
                        help="Output video FPS")

    # --- torch.compile ---
    parser.add_argument("--compile", action="store_true",
                        help="Wrap diffusion model + Wan VAE encoder/decoder + "
                             "TAEHV (when present) with torch.compile. First "
                             "warmup clip absorbs the compile time; subsequent "
                             "clips run on the compiled graphs.")

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


# ===========================================================================
# Inference & post-processing
# ===========================================================================

@torch.no_grad()
def run_inference(
    model, condition, num_latent_frames, t_list,
    chunk_size, context_noise, seed, device, dtype,
):
    """Block-wise AR inference loop.

    Adapted from Self-Forcing's rollout_with_gradient but inference-only:
    - No gradients, no random exit steps
    - Full denoising per block (all steps in t_list)
    - KV cache updated after each block with denoised output
    - Rolling window eviction handled internally by CausalSelfAttention

    Args:
        model: CausalOmniAvatarWan (1.3B student)
        condition: dict with text_embeds, audio_emb, ref_latent, mask, etc.
        num_latent_frames: total latent frames to generate
        t_list: denoising timestep schedule (e.g. [0.999, 0.9, 0.75, 0.5, 0.0])
        chunk_size: frames per AR block (3)
        context_noise: noise level for cache updates (0 = clean)
        seed: random seed
        device: torch device
        dtype: torch dtype

    Returns:
        output: [1, 16, num_latent_frames, H_lat, W_lat] denoised latents
    """
    # Update model's total_num_frames for correct cache allocation
    model.total_num_frames = num_latent_frames
    model.clear_caches()

    # Determine spatial dims from ref_latent
    ref_latent = condition["ref_latent"]  # [1, 16, 1, H_lat, W_lat]
    B = ref_latent.shape[0]
    C = 16
    H_lat, W_lat = ref_latent.shape[3], ref_latent.shape[4]

    num_blocks = num_latent_frames // chunk_size
    assert num_latent_frames % chunk_size == 0

    # Generate noise
    torch.manual_seed(seed)
    noise = torch.randn(B, C, num_latent_frames, H_lat, W_lat, device=device, dtype=dtype)

    # Convert t_list to tensor
    t_list_t = torch.tensor(t_list, device=device, dtype=torch.float64)

    # Output accumulator
    output = torch.zeros_like(noise)

    print(f"  {num_blocks} blocks x {len(t_list) - 1} denoising steps")
    for block_idx in range(num_blocks):
        cur_start_frame = block_idx * chunk_size

        # Slice noise for this chunk
        noisy_input = noise[:, :, cur_start_frame:cur_start_frame + chunk_size]

        # Multi-step denoising
        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]

            # Forward pass — model.forward() handles _build_y, rescale_t, _forward_ar internally
            # Keep timesteps in float64 for numerically stable scheduling
            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:
                # Add noise for next step (SDE: fresh random noise)
                eps = torch.randn_like(x0_pred)
                noisy_input = model.noise_scheduler.forward_process(
                    x0_pred, eps, t_next.expand(B),
                )
            else:
                # Final step — clean output
                noisy_input = x0_pred

        # Store denoised chunk
        output[:, :, cur_start_frame:cur_start_frame + chunk_size] = x0_pred

        # Update KV cache with denoised output (context for next block)
        cache_input = x0_pred
        t_cache = torch.full((B,), context_noise, device=device, dtype=torch.float64)
        if context_noise > 0:
            cache_eps = torch.randn_like(x0_pred)
            cache_input = model.noise_scheduler.forward_process(
                x0_pred, cache_eps,
                torch.tensor(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) % 10 == 0 or block_idx == num_blocks - 1:
            print(f"  Block {block_idx + 1}/{num_blocks} done")

    model.clear_caches()
    return output


@torch.no_grad()
def decode_and_save(vae, output_latents, audio_path, output_path, fps, device):
    """VAE decode latents -> save silent video -> mux with audio."""
    import imageio.v3 as iio

    # VAE decode — expects list of [C, T_lat, H_lat, W_lat] in float32
    latent_for_vae = output_latents[0].to(torch.float32)  # [16, T_lat, H_lat, W_lat]
    video_tensor = vae.decode([latent_for_vae], device=device)  # [1, 3, T_video, H, W]
    video_tensor = video_tensor.clamp(-1, 1)

    # Convert to uint8 frames: [T, H, W, 3]
    video_np = video_tensor[0]  # [3, T, H, W]
    video_np = video_np.permute(1, 2, 3, 0)  # [T, H, W, 3]
    video_np = ((video_np.float() + 1) * 127.5).clamp(0, 255).cpu().to(torch.uint8).numpy()

    # Save silent video to temp file
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    tmp_silent = output_path + ".silent.mp4"
    iio.imwrite(
        tmp_silent,
        video_np,
        fps=fps,
        codec="libx264",
        output_params=["-loglevel", "quiet", "-crf", "18"],
    )
    print(f"  Silent video: {video_np.shape[0]} frames at {fps}fps")

    # Mux with audio
    video_duration = video_np.shape[0] / fps
    mux_video_with_audio(tmp_silent, audio_path, output_path, duration_s=video_duration)

    # Cleanup
    if os.path.exists(tmp_silent):
        os.remove(tmp_silent)


# ===========================================================================
# Main
# ===========================================================================

def main():
    args = parse_args()
    validate_args(args)
    use_preprocessing = not args.skip_preprocessing

    # Activate per-function torch.compile decorators in network_causal.py
    # BEFORE the model class is imported (which happens later via
    # load_diffusion_model). Done by setting LIPFORCING_COMPILE=true.
    if args.compile:
        os.environ["LIPFORCING_COMPILE"] = "true"

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

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

    # ===================================================================
    # Load models once (expensive — minutes for 14B/1.3B weights)
    # ===================================================================
    print("Loading diffusion model ...")
    model = load_diffusion_model(args, device, dtype)

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

    # Decoder selection: full Wan VAE stays loaded for encoding the driving video;
    # decoding swaps to TAEHV tiny decoder if --taehv_ckpt is provided.
    if args.taehv_streaming:
        if not args.taehv_ckpt:
            raise ValueError("--taehv_streaming requires --taehv_ckpt")
        print(f"Loading StreamingTAEHV decoder from {args.taehv_ckpt} ...")
        decoder_vae = StreamingTAEHVDecoderWrapper(args.taehv_ckpt, device)
    elif args.taehv_ckpt:
        print(f"Loading TAEHV tiny decoder from {args.taehv_ckpt} ...")
        decoder_vae = TAEHVDecoderWrapper(args.taehv_ckpt, device)
    else:
        decoder_vae = vae

    # Encoder selection: default to full Wan VAE. If --taehv_encode is set,
    # reuse the same TAEHV model (it implements both encode and decode).
    if args.taehv_encode:
        if not args.taehv_ckpt:
            raise ValueError("--taehv_encode requires --taehv_ckpt")
        print("Using TAEHV tiny encoder for driving video encoding.")
        encoder_vae = decoder_vae
    else:
        encoder_vae = vae

    # Eagerly load Wav2Vec + text
    wav2vec_model = wav2vec_extractor = None
    if args.wav2vec_path:
        print("Loading Wav2Vec2 (eager) ...")
        wav2vec_model, wav2vec_extractor = load_wav2vec(args.wav2vec_path, device)
        # Warmup forward pass to compile CUDA kernels.
        # OmniAvatar Wav2VecModel requires seq_len + output_hidden_states.
        _dummy_audio = np.zeros(16000, dtype=np.float32)  # 1s @ 16kHz → 25 video-frames
        _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 (eager) ...")
    text_embeds = load_or_encode_text(args, device, dtype)

    # LatentSync ImageProcessor (face detection + alignment; on by default)
    image_processor = None
    if use_preprocessing:
        image_processor = load_image_processor(args.mask_path, device)

    # ===================================================================
    # Optional torch.compile wrapping (compile time absorbed by warmup)
    # ===================================================================
    # Hot DiT functions are decorated via @conditional_compile (activated by
    # LIPFORCING_COMPILE=true env var, set above before model imports).  Here
    # we additionally wrap the VAE / TAEHV encode + decode forwards.
    if args.compile:
        print("[--compile] Hot DiT functions decorated with @conditional_compile. "
              "Warmup clip will trigger Dynamo trace.")
        _compile_kw = dict(mode=None, backend="inductor", dynamic=None)

        # Compile VAE encode/decode paths.
        # TAEHV is skipped — its internals do `b = model[i]` which breaks
        # when the Sequential is wrapped in an OptimizedModule.  The DiT
        # gets compile via @conditional_compile (LIPFORCING_COMPILE=true above).

        # Wan VAE decoder compile (skip if decoder is TAEHV)
        if not isinstance(decoder_vae, (TAEHVDecoderWrapper, StreamingTAEHVDecoderWrapper)):
            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 (skip if encoder is TAEHV)
        if not isinstance(encoder_vae, (TAEHVDecoderWrapper, StreamingTAEHVDecoderWrapper)):
            if hasattr(encoder_vae, 'model') and hasattr(encoder_vae.model, 'encoder'):
                if not isinstance(encoder_vae.model.encoder, torch._dynamo.eval_frame.OptimizedModule):
                    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}")

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

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

        tmp_audio = None
        try:
            # --- Resolve audio ---
            audio_path, tmp_audio = resolve_audio(
                audio_path=audio_path_sample, video_path=video_path,
            )

            # --- Compute generation length ---
            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,
            )

            # --- Face detection + alignment (default preprocessing) ---
            latentsync_metadata = None
            if use_preprocessing:
                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(f"  [FAIL] LatentSync preprocessing failed, skipping {name}")
                    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,
                )
            else:
                # Wav2Vec + text already loaded eagerly before the loop

                # Reference frames: aligned faces from LatentSync or raw video
                if use_preprocessing and latentsync_metadata is not None:
                    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)
                else:
                    ref_frames_np = load_and_adjust_video(video_path, num_video_frames)

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

            # --- Run inference ---
            print("Running inference ...")
            output_latents = run_inference(
                model, condition, num_latent_frames, args.t_list,
                args.chunk_size, args.context_noise, args.seed, device, dtype,
            )

            # --- Post-processing: decode + save ---
            os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

            if use_preprocessing and latentsync_metadata is not None:
                # LatentSync compositing path — float-space decode + composite
                print("VAE decoding (float) ...")
                latent_for_vae = output_latents[0].to(torch.float32)
                video_decoded = decoder_vae.decode([latent_for_vae], device=device)
                video_decoded = video_decoded.clamp(-1, 1)

                # [1, 3, T_video, H, W] -> [T, 3, H, W] in [0, 1]
                generated_float = video_decoded[0].permute(1, 0, 2, 3)  # [3,T,H,W] -> [T,3,H,W]
                generated_float = ((generated_float + 1) / 2).clamp(0, 1)  # [-1,1] -> [0,1]

                # Composite onto original frames
                print("Compositing ...")
                composited_np = composite_with_latentsync_float(
                    generated_float.cpu(), latentsync_metadata, image_processor,
                    use_mouth_only_compositing=not args.composite_full_face,
                )

                # Save composited video (original resolution) with audio
                composited_path = output_path
                save_frames_as_video(composited_np, composited_path, fps=args.fps)
                video_duration = composited_np.shape[0] / args.fps
                tmp_composited = composited_path + ".tmp.mp4"
                os.rename(composited_path, tmp_composited)
                mux_video_with_audio(tmp_composited, audio_path, composited_path,
                                     duration_s=video_duration)
                if os.path.exists(tmp_composited):
                    os.remove(tmp_composited)

                print(f"  Saved composited: {composited_path}")

                # Optionally also save the raw generated aligned (512x512) video
                if args.save_aligned:
                    aligned_path = output_path.replace(".mp4", "_aligned.mp4")
                    aligned_np = ((generated_float.permute(0, 2, 3, 1).cpu().float()) * 255
                                  ).clamp(0, 255).to(torch.uint8).numpy()
                    save_frames_as_video(aligned_np, aligned_path, fps=args.fps)
                    tmp_aligned = aligned_path + ".tmp.mp4"
                    os.rename(aligned_path, tmp_aligned)
                    mux_video_with_audio(tmp_aligned, audio_path, aligned_path,
                                         duration_s=video_duration)
                    if os.path.exists(tmp_aligned):
                        os.remove(tmp_aligned)
                    print(f"  Saved aligned:    {aligned_path}")
            else:
                # Standard decode + save (no LatentSync)
                print("Decoding and saving ...")
                decode_and_save(decoder_vae, output_latents, audio_path, output_path,
                                args.fps, device)

            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:
            # Cleanup per-sample temp audio
            if tmp_audio is not None and os.path.exists(tmp_audio):
                os.remove(tmp_audio)

            # Free per-sample GPU memory
            torch.cuda.empty_cache()

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



if __name__ == "__main__":
    main()