eval_frame.py

import argparse
import csv
import json
import os
import random
import re
import shutil
import sys
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from datetime import datetime
from glob import glob
from pathlib import Path
from typing import List

import numpy as np
import torch
import torchaudio
from lightning.pytorch import seed_everything
from tqdm import tqdm

from prefigure.prefigure import get_all_args, push_wandb_config

from ThinkSound.data.datamodule import DataModule
from ThinkSound.inference.sampling import sample, sample_discrete_euler
from ThinkSound.models import create_model_from_config
from ThinkSound.models.diffusion import MMDiTWrapper, MMTPDiTWrapper
from ThinkSound.models.utils import load_ckpt_state_dict, remove_weight_norm_from_model
from detect_all import CANON_LABELS, _labels_suffix, _parse_ih_labels


def load_checkpoint(model, ckpt_path, prefix="diffusion."):
    print(f"Loading checkpoint from: {ckpt_path}")
    try:
        ckpt = torch.load(ckpt_path, map_location="cpu")
    except Exception as e:
        print(f"Error loading checkpoint file: {e}")
        return

    state_dict = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
    processed_state_dict = {}

    for k, v in state_dict.items():
        if k.startswith(prefix):
            new_key = k[len(prefix):]
            processed_state_dict[new_key] = v
        elif k.startswith("diffusion_ema."):
            continue
        else:
            processed_state_dict[k] = v

    try:
        model.load_state_dict(processed_state_dict)
        print("Checkpoint loaded successfully!")
    except RuntimeError as e:
        print(f"Error loading state_dict: {e}")
        print("Attempting to load with strict=False...")
        model.load_state_dict(processed_state_dict, strict=False)
        print("Checkpoint loaded with strict=False. Some keys may be missing or unexpected.")


@torch.no_grad()
def predict_step(diffusion, batch, diffusion_objective, use_video_caption=False, device='cuda:0'):
    diffusion = diffusion.to(device)
    reals, metadata = batch
    ids = [item['id'] for item in metadata]
    batch_size, length = reals.shape[0], reals.shape[2]
    
    with torch.amp.autocast('cuda'):
        conditioning = diffusion.conditioner(metadata, device)

    if not use_video_caption:
        conditioning['t5_features'] = conditioning.pop('t5_features_caption')
    else:
        del conditioning["t5_features_caption"]

    video_exist = torch.stack([item['video_exist'] for item in metadata], dim=0)
    conditioning['metaclip_features'][~video_exist] = diffusion.model.model.empty_clip_feat
    conditioning['sync_features'][~video_exist] = diffusion.model.model.empty_sync_feat

    cond_inputs = diffusion.get_conditioning_inputs(conditioning)
    if isinstance(diffusion.model, MMTPDiTWrapper):
        cond_inputs.update({'timing_plan': [data['timing_plan'] for data in metadata]})
        cond_inputs.update({'event_t5_features': [data['event_t5_features'] for data in metadata]})

    if batch_size > 1:
        noise = torch.cat([torch.randn([1, diffusion.io_channels, length]).to(device) for _ in range(batch_size)], dim=0)
    else:
        noise = torch.randn([batch_size, diffusion.io_channels, length]).to(device)

    with torch.amp.autocast('cuda'):
        model = diffusion.model
        if diffusion_objective == "v":
            fakes = sample(model, noise, 24, 0, **cond_inputs, cfg_scale=5, batch_cfg=True)
        elif diffusion_objective == "rectified_flow":
            fakes = sample_discrete_euler(model, noise, 24, **cond_inputs, cfg_scale=5, batch_cfg=True)
            
        if diffusion.pretransform is not None:
            fakes = diffusion.pretransform.decode(fakes)

    audios = fakes.to(torch.float32).div(torch.max(torch.abs(fakes))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
    return audios


def copy_non_ih_wavs(all_wavs, ih_basenames, audio_dir, move=False, max_workers=8):
    os.makedirs(audio_dir, exist_ok=True)

    def copy_one(src):
        base = os.path.basename(src)
        if base in ih_basenames:  # Skip if in hallucination list
            return False
        dst = os.path.join(audio_dir, base)
        if os.path.exists(dst):   # Skip if already copied
            return False
        shutil.copy2(src, dst)
        return True

    moved = 0
    with ThreadPoolExecutor(max_workers=max_workers) as ex:
        for ok in ex.map(copy_one, all_wavs):
            if ok:
                moved += 1

    return moved


def collect_ih_and_copy_others(filled_first_root_dir: str, audio_dir: str, json_name: str) -> List[str]:
    """
    Reads json_name from filled_first_root_dir/cache,
    collects wavs with IH@vid != 0 (hallucinated),
    and copies the rest to audio_dir.
    """
    cache_json = os.path.join(filled_first_root_dir, "cache", json_name)
    if not os.path.exists(cache_json):
        raise FileNotFoundError(f"JSON file not found: {cache_json}")

    with open(cache_json, "r", encoding="utf-8") as f:
        data = json.load(f)

    try:
        mv_metrics = data["Individual Audio Metrics"]
    except KeyError as e:
        raise KeyError("Missing ['Individual Audio Metrics'] in JSON") from e

    # Collect filenames with hallucinations
    ih_names = [name for name, vals in mv_metrics.items() if float(vals.get("IH@vid", 0)) != 0]

    ih_wavs: List[str] = []
    missing: List[str] = []
    for name in ih_names:
        full = os.path.join(filled_first_root_dir, name)
        if os.path.exists(full):
            ih_wavs.append(full)
        else:
            missing.append(name)

    # Copy non-hallucinated wavs to audio_dir
    os.makedirs(audio_dir, exist_ok=True)
    ih_basenames = {os.path.basename(p) for p in ih_wavs}

    all_wavs = [
        os.path.join(filled_first_root_dir, fn)
        for fn in os.listdir(filled_first_root_dir)
        if fn.lower().endswith(".wav")
    ]

    moved = copy_non_ih_wavs(
        all_wavs=all_wavs,
        ih_basenames=ih_basenames,
        audio_dir=audio_dir,
        max_workers=8
    )

    if missing:
        print(f"[WARN] {len(missing)} files from JSON not found in dir, e.g.: {missing[:5]}")

    return ih_wavs 


def load_segments_from_csv(csv_path, ih_labels=("speech","music")):
    """
    Reads CSV and returns a dictionary mapped by filename, containing segment info.
    """
    segments = {}
    
    with open(csv_path, 'r', newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            filename = row['filename'].strip().replace('.wav', '').replace('.flac', '')
            start = float(row['start'])
            end = float(row['end'])
            label = row['label'].strip()

            if label not in ih_labels:
                continue

            if filename not in segments:
                segments[filename] = []

            segments[filename].append({
                'start': start,
                'end': end,
                'label': label
            })
    
    return segments


def replace_features(
    conditioning,
    diffusion_model,
    ids,
    csv_path,
    args,
    retri_f,
    ih_labels=("speech","music"),
    clip_fps=8,
    sync_fps=24,
    replace_metaclip=False,
    replace_sync=False,
    reverse_replace=False,
    zero_replace=False,
    random_replace=False,
    random_p=0.6035,
):
    """
    Replaces metaclip/sync features.
    - random_replace=True: Replaces randomly with probability p.
    - random_replace=False:
        reverse_replace=False: Masks (start, end) segments.
        reverse_replace=True: Masks outside (start, end) segments.
    - zero_replace=True: Replaces with zero vectors, otherwise empty_feat.
    """
    segments = load_segments_from_csv(csv_path, ih_labels)

    metaclip_features = conditioning['metaclip_features']
    sync_features = conditioning['sync_features']

    def make_repl(empty_feat, like_tensor):
        if zero_replace:
            return torch.zeros_like(empty_feat, device=like_tensor.device, dtype=like_tensor.dtype)
        else:
            return empty_feat.to(device=like_tensor.device, dtype=like_tensor.dtype)

    empty_clip_1 = diffusion_model.model.empty_clip_feat
    empty_sync_1 = diffusion_model.model.empty_sync_feat
    if args.retri_replace:
        empty_clip_1 = retri_f 

    def _bounded_range(T, s, e):
        s = max(0, min(int(s), T))
        e = max(0, min(int(e), T))
        if e < s:
            e = s
        return s, e

    def _fill_mask_with_adjacent(x_bt_d: torch.Tensor, mask_t: torch.Tensor, fallback_vec_1d: torch.Tensor):
        T, D = x_bt_d.shape
        if mask_t.sum().item() == 0:
            return

        fb = fallback_vec_1d
        if fb.dim() == 2:
            fb = fb[0]
        fb = fb.to(device=x_bt_d.device, dtype=x_bt_d.dtype)

        m = mask_t.bool()
        idx = torch.nonzero(m, as_tuple=False).flatten()
        if idx.numel() == 0:
            return

        # Scan for continuous segments
        runs = []
        start = idx[0].item()
        prev = start
        for i in idx[1:].tolist():
            if i == prev + 1:
                prev = i
            else:
                runs.append((start, prev + 1))
                start = i
                prev = i
        runs.append((start, prev + 1))

        # Fill segments using adjacent frames
        for s, e in runs:
            left = s - 1
            right = e

            use_vec = None
            if left >= 0 and (not m[left].item()):
                use_vec = x_bt_d[left]
            elif right < T and (not m[right].item()):
                use_vec = x_bt_d[right]
            else:
                use_vec = fb

            x_bt_d[s:e, :] = use_vec.unsqueeze(0).expand(e - s, -1)


    for b_idx, sample_id in enumerate(ids):
        if sample_id not in segments:
            continue
            
        # MetaClip processing
        if replace_metaclip and metaclip_features is not None:
            T1 = metaclip_features.shape[1]
            repl_1 = make_repl(empty_clip_1, metaclip_features[b_idx])
            mask_clip = torch.zeros(T1, dtype=torch.bool, device=metaclip_features.device)

            if random_replace:
                mask_clip = (torch.rand(T1, device=metaclip_features.device) < float(random_p))
            else:
                if sample_id in segments:
                    segs = segments[sample_id]
                    if reverse_replace:
                        pos_mask = torch.zeros(T1, dtype=torch.bool, device=metaclip_features.device)
                        for seg in segs:
                            s, e = _bounded_range(T1, seg['start'] * clip_fps, seg['end'] * clip_fps)
                            pos_mask[s:e] = True
                        mask_clip = ~pos_mask
                    else:
                        for seg in segs:
                            s, e = _bounded_range(T1, seg['start'] * clip_fps, seg['end'] * clip_fps)
                            mask_clip[s:e] = True

            n = mask_clip.sum().item()
            if n > 0:
                if args.adj_replace:
                    _fill_mask_with_adjacent(metaclip_features[b_idx], mask_clip, repl_1)
                else:
                    metaclip_features[b_idx, mask_clip, :] = repl_1.expand(n, -1)

        # Sync processing
        if replace_sync and sync_features is not None:
            T2 = sync_features.shape[1]
            repl_1 = make_repl(empty_sync_1, sync_features[b_idx])
            mask_sync = torch.zeros(T2, dtype=torch.bool, device=sync_features.device)

            if random_replace:
                mask_sync = (torch.rand(T2, device=sync_features.device) < float(random_p))
            else:
                if sample_id in segments:
                    segs = segments[sample_id]
                    if reverse_replace:
                        pos_mask = torch.zeros(T2, dtype=torch.bool, device=sync_features.device)
                        for seg in segs:
                            s, e = _bounded_range(T2, seg['start'] * sync_fps, seg['end'] * sync_fps)
                            pos_mask[s:e] = True
                        mask_sync = ~pos_mask
                    else:
                        for seg in segs:
                            s, e = _bounded_range(T2, seg['start'] * sync_fps, seg['end'] * sync_fps)
                            mask_sync[s:e] = True

            n = mask_sync.sum().item()
            if n > 0:
                sync_features[b_idx, mask_sync, :] = repl_1.expand(n, -1)

    return conditioning


def predict_step_replace(
    diffusion, batch, diffusion_objective, duration, csv_path, args, 
    retri_f=None, ih_labels=("speech","music"), replace_metaclip=False, 
    replace_sync=False, reverse_replace=False, zero_replace=False, 
    random_replace=False, use_video_caption=False, device='cuda:0'
):
    diffusion = diffusion.to(device)

    reals, metadata = batch
    ids = [item['id'] for item in metadata]
    batch_size, length = reals.shape[0], reals.shape[2]
    
    with torch.amp.autocast('cuda'):
        conditioning = diffusion.conditioner(metadata, device)
    
    if not use_video_caption:
        conditioning['t5_features'] = conditioning.pop('t5_features_caption') # simple label text
    else:
        del conditioning["t5_features_caption"] # rich semantic caption

    video_exist = torch.stack([item['video_exist'] for item in metadata], dim=0)
    conditioning['metaclip_features'][~video_exist] = diffusion.model.model.empty_clip_feat
    conditioning['sync_features'][~video_exist] = diffusion.model.model.empty_sync_feat
    
    conditioning = replace_features(
        conditioning, diffusion.model, ids, csv_path, args=args, retri_f=retri_f, 
        ih_labels=ih_labels, replace_metaclip=replace_metaclip, replace_sync=replace_sync, 
        reverse_replace=reverse_replace, zero_replace=zero_replace, random_replace=random_replace
    )

    cond_inputs = diffusion.get_conditioning_inputs(conditioning)
    if isinstance(diffusion.model, MMTPDiTWrapper):
        cond_inputs.update({'timing_plan': [data['timing_plan'] for data in metadata]})
        cond_inputs.update({'event_t5_features': [data['event_t5_features'] for data in metadata]})
    
    if batch_size > 1:
        noise_list = []
        for _ in range(batch_size):
            noise_1 = torch.randn([1, diffusion.io_channels, length]).to(device)  # Advance RNG state
            noise_list.append(noise_1)
        noise = torch.cat(noise_list, dim=0)
    else:
        noise = torch.randn([batch_size, diffusion.io_channels, length]).to(device)

    with torch.amp.autocast('cuda'):
        model = diffusion.model
        if diffusion_objective == "v":
            fakes = sample(model, noise, 24, 0, **cond_inputs, cfg_scale=5, batch_cfg=True)
        elif diffusion_objective == "rectified_flow":
            fakes = sample_discrete_euler(model, noise, 24, **cond_inputs, cfg_scale=5, batch_cfg=True)
            
        if diffusion.pretransform is not None:
            fakes = diffusion.pretransform.decode(fakes)

    audios = fakes.to(torch.float32).div(torch.max(torch.abs(fakes))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
    return audios


def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def get_all_args_with_custom():
    pre_parser = argparse.ArgumentParser(add_help=False)
    pre_parser.add_argument("--debug", action="store_true")
    pre_parser.add_argument("--use_video_caption", action="store_true", help="Whether to use video caption")
    pre_parser.add_argument("--root-dir", type=str, default="", help="Root dir for L1/L2 enumeration. Inferred from --results-dir if empty.")
    pre_parser.add_argument("--first-root-dir", type=str, default="")
    pre_parser.add_argument("--csv-path", type=str, default=None)
    pre_parser.add_argument("--ih_labels", default="speech,music", help="Comma-separated labels, e.g., speech,music or bird,water")
    pre_parser.add_argument("--replace_metaclip", action="store_true")
    pre_parser.add_argument("--replace_sync", action="store_true")
    pre_parser.add_argument("--reverse_replace", action="store_true")
    pre_parser.add_argument("--zero_replace", action="store_true")
    pre_parser.add_argument("--random_replace", action="store_true")
    pre_parser.add_argument("--adj_replace", action="store_true")
    pre_parser.add_argument("--retri_replace", action="store_true")
    
    custom_args, remaining_argv = pre_parser.parse_known_args(sys.argv[1:])
    _orig_argv = sys.argv[:]
    
    try:
        sys.argv = [_orig_argv[0]] + remaining_argv
        base_args = get_all_args()
    finally:
        sys.argv = _orig_argv

    for k, v in vars(custom_args).items():
        setattr(base_args, k, v)

    if base_args.replace_metaclip or base_args.replace_sync:
        if base_args.reverse_replace:
            base_args.save_dir = base_args.save_dir.replace("thinksound", "re_thinksound_replace")
        elif base_args.zero_replace:
            base_args.save_dir = base_args.save_dir.replace("thinksound", "zero_thinksound_replace")
        elif base_args.random_replace:
            base_args.save_dir = base_args.save_dir.replace("thinksound", "random_thinksound_replace")
        else:
            base_args.save_dir = base_args.save_dir.replace("thinksound", "thinksound_replace")
        
    if base_args.replace_metaclip and not base_args.replace_sync:
        base_args.save_dir = base_args.save_dir.replace("thinksound_replace", "thinksound_replace_metaclip")
    elif base_args.replace_sync and not base_args.replace_metaclip:
        base_args.save_dir = base_args.save_dir.replace("thinksound_replace", "thinksound_replace_sync")

    print(f"use_video_caption: {base_args.use_video_caption}")
    print(f"root_dir: {base_args.root_dir or '(auto)'}")
    
    return base_args


def infer_root_dir_from_template(path_with_placeholders: str) -> str:
    """
    Infers root directory from a template containing {L1_path}.
    """
    if not path_with_placeholders:
        return ""
    tok = "{L1_path}"
    if tok in path_with_placeholders:
        return path_with_placeholders.split(tok)[0].rstrip("/")

    # Fallback if no placeholders: return two levels up
    p = Path(path_with_placeholders)
    return str(p.parent.parent)


def list_l1_l2_pairs(root_dir: str):
    """
    Scans two levels under root_dir and returns a sorted list of (L1, L2) pairs.
    """
    pairs = []
    root = Path(root_dir)
    
    if not root.exists():
        print(f"[WARN] root_dir does not exist: {root_dir}")
        return pairs
        
    for l1 in sorted([d for d in root.iterdir() if d.is_dir()]):
        for l2 in sorted([d for d in l1.iterdir() if d.is_dir()]):
            pairs.append((l1.name, l2.name))
            
    return pairs


def fill_placeholders(s: str, L1: str, L2: str) -> str:
    if s is None:
        return s
    s = s.replace("{L1_path}", L1)
    s = s.replace("{L2_path}", L2)
    return s


def main():
    args = get_all_args_with_custom()

    if args.save_dir == '':
        args.save_dir = args.results_dir

    # Infer root_dir
    root_dir = args.root_dir.strip()
    if not root_dir:
        root_dir = args.results_dir
        if not root_dir and args.dataset_config:
            try:
                with open(args.dataset_config) as f:
                    tmp_cfg = json.load(f)
                if tmp_cfg["test_datasets"]:
                    root_dir = infer_root_dir_from_template(tmp_cfg["test_datasets"][0].get("path", ""))
            except Exception:
                pass
                
    print(f"[INFO] Using root_dir = {root_dir}")

    # Set random seed
    seed = args.seed
    if os.environ.get("SLURM_PROCID") is not None:
        seed += int(os.environ.get("SLURM_PROCID"))
    seed_everything(seed, workers=True)

    # Load model config and model
    if args.model_config == '':
        args.model_config = "ThinkSound/configs/model_configs/thinksound.json"
    print("Use AR Model" if '_ar' in args.model_config else "Use Common Model")

    with open(args.model_config) as f:
        model_config = json.load(f)

    duration = float(args.duration_sec)
    sample_rate = model_config["sample_rate"]
    latent_length = int(44100 / 64 / 32 * duration)

    model_config["sample_size"] = duration * sample_rate
    Tin = int(25 * duration)
    sync_seq_len = ((Tin - 16) // 8 + 1) * 8
    model_config["model"]["diffusion"]["config"]["sync_seq_len"] = sync_seq_len
    model_config["model"]["diffusion"]["config"]["clip_seq_len"] = 8 * int(duration)
    model_config["model"]["diffusion"]["config"]["latent_seq_len"] = latent_length

    if not args.debug:
        model = create_model_from_config(model_config)
        if args.compile:
            model = torch.compile(model)

        load_checkpoint(model, args.ckpt_dir, prefix="diffusion.")
        vae_state = load_ckpt_state_dict(args.pretransform_ckpt_path, prefix='autoencoder.')
        model.pretransform.load_state_dict(vae_state)

    if args.dataset_config == '':
        args.dataset_config = "ThinkSound/configs/multimodal_dataset_demo.json"
    with open(args.dataset_config) as f:
        dataset_config_template = json.load(f)

    # Iterate over L1/L2 pairs
    pairs = list_l1_l2_pairs(root_dir)
    if not pairs:
        print("[WARN] No L1/L2 pairs found. Please verify root_dir structure is root/L1/L2.")
        return

    current_date = datetime.now()
    formatted_date = current_date.strftime('%m%d')

    for L1, L2 in pairs:
        filled_results_dir = fill_placeholders(args.results_dir, L1, L2)
        filled_save_dir = fill_placeholders(args.save_dir, L1, L2)
        filled_csv_path = fill_placeholders(args.csv_path, L1, L2)
        
        audio_dir = filled_save_dir
        
        # Check if output already exists
        if (os.path.exists(audio_dir) 
            and glob(os.path.join(audio_dir, "*.wav")) 
            and not args.debug):
            print(f"[SKIP] .wav files exist, skipping L1={L1}, L2={L2} -> {filled_save_dir}")
            continue
            
        os.makedirs(audio_dir, exist_ok=True)

        dataset_config = deepcopy(dataset_config_template)
        for td in dataset_config.get("test_datasets", []):
            if "path" in td and isinstance(td["path"], str):
                td["path"] = fill_placeholders(td["path"], L1, L2)
            else:
                td["path"] = filled_results_dir

            split_default = str(Path(root_dir) / L1 / L2 / "video.txt")
            if "split_path" in td and isinstance(td["split_path"], str) and td["split_path"]:
                td["split_path"] = fill_placeholders(td["split_path"], L1, L2)
            else:
                td["split_path"] = split_default

        # Validate existence
        td0 = dataset_config["test_datasets"][0]
        data_path = Path(td0["path"])
        split_path = Path(td0["split_path"])
        
        if not data_path.exists():
            print(f"[SKIP] Data directory not found: {data_path} (L1={L1}, L2={L2})")
            continue
        if not split_path.exists():
            print(f"[SKIP] split file not found: {split_path} (L1={L1}, L2={L2})")
            continue

        print(f"\n[RUN] L1={L1}  L2={L2}")
        print(f"      data_path  = {data_path}")
        print(f"      split_path = {split_path}")

        dm = DataModule(
            dataset_config,
            batch_size=args.batch_size,
            test_batch_size=args.test_batch_size,
            num_workers=args.num_workers,
            sample_rate=sample_rate,
            sample_size=float(args.duration_sec) * sample_rate,
            audio_channels=model_config.get("audio_channels", 2),
            latent_length=latent_length,
        )
        
        ih_labels = _parse_ih_labels(getattr(args, "ih_labels", "speech,music"))
        ih_suf = _labels_suffix(ih_labels)
        json_name = f"ih_metrics_panns_{ih_suf}.json"
        
        if args.replace_metaclip or args.replace_sync:
            filled_first_root_dir = fill_placeholders(args.first_root_dir, L1, L2)
            ih_list = collect_ih_and_copy_others(filled_first_root_dir, audio_dir, json_name)
            dm.setup('predict', ih_list=ih_list)
        else:
            dm.setup('predict')
            
        dl = dm.predict_dataloader()

        if args.retri_replace:
            retri_path = filled_csv_path.replace('ThinkSound', 'KlingFoley').replace('cache/segment_fused_or.csv', 'avg_metaclip_proto.pth')
            data = torch.load(retri_path, map_location="cuda")
            retri_f = data['proto']
        else:
            retri_f = None

        # Prediction loop
        for batch in tqdm(dl, desc=f"Predicting [{L1}/{L2}]"):
            if args.replace_metaclip or args.replace_sync:
                audio = predict_step_replace(
                    model,
                    batch=batch,
                    diffusion_objective=model_config["model"]["diffusion"]["diffusion_objective"],
                    duration=duration,
                    csv_path=filled_csv_path,
                    args=args,
                    retri_f=retri_f,
                    ih_labels=ih_labels,
                    replace_metaclip=args.replace_metaclip, 
                    replace_sync=args.replace_sync,
                    reverse_replace=args.reverse_replace,
                    zero_replace=args.zero_replace,
                    random_replace=args.random_replace,
                    use_video_caption=args.use_video_caption,
                    device='cuda:0'
                )
            else:
                audio = predict_step(
                    model,
                    batch=batch,
                    diffusion_objective=model_config["model"]["diffusion"]["diffusion_objective"],
                    use_video_caption=args.use_video_caption,
                    device='cuda:0'
                )

            _, metadata = batch
            ids = [item['id'] for item in metadata]

            for i in range(audio.size(0)):
                id_str = ids[i] if i < len(ids) else f"unknown_{i}"
                out_path = os.path.join(audio_dir, f"{id_str}.wav")
                torchaudio.save(out_path, audio[i], 44100)
                
        print(f"[DONE] Output directory: {audio_dir}")

    print("\n[ALL DONE] All L1/L2 pairs processed.")


if __name__ == '__main__':
    main()