diff --git a/datasets/README.md b/datasets/README.md index 510d15b68..04fdf3f01 100644 --- a/datasets/README.md +++ b/datasets/README.md @@ -63,6 +63,7 @@ needs a `.gitignore` line, which keeps that decision explicit. | Folder | What | |---|---| +| [`audiocaps/`](audiocaps/README.md) | AudioCaps captions for LTX-2.3 CLAP audio-reward RL | | [`arxivqa_mc/`](arxivqa_mc/README.md) | ArxivQA scientific-figure multiple-choice (BAGEL GRPO) | | [`asearcher/`](asearcher/README.md) | ASearcher deep-research prompts (agentic RL) | | [`daily_omni_av/`](daily_omni_av/README.md) | Daily-Omni audio-video QA | diff --git a/datasets/audiocaps/README.md b/datasets/audiocaps/README.md new file mode 100644 index 000000000..2012b8ffe --- /dev/null +++ b/datasets/audiocaps/README.md @@ -0,0 +1,95 @@ +# AudioCaps for LTX-2.3 CLAP training + +This converter prepares the official human-written captions from +[AudioCaps](https://github.com/cdjkim/audiocaps) for the LTX-2.3 CLAP reward recipe. +AudioCaps was introduced at NAACL 2019 and contains captions for AudioSet clips. + +The generated JSONL is a local artifact and must not be committed. This recipe does not +download or consume the source audio: LTX-2.3 generates audio from each caption and CLAP +scores that generated waveform against the same official caption. + +The event-grounded variant also joins the AudioCaps clips to their human-confirmed AudioSet +labels. It uses those labels as positive sound-event targets and RMS-normalizes generated +audio before CLAP scoring so increasing waveform gain cannot improve the reward by itself. + +## Source and terms + +By default, the converter reads the official CSV files at AudioCaps commit +`d004db3ea1b01cf4fd0347dd8d27db90cadc8809`: + +- train: 49,838 clips with one caption each +- validation: 495 clips with five captions each (2,475 caption rows) +- test: 975 clips with five captions each (4,875 caption rows) + +The upstream repository says its code and dataset are free to use for academic purposes and +asks users to cite the AudioCaps paper. Review the upstream terms before redistributing or +using the data outside that scope. + +## Cook + +From the repository root: + +```bash +python datasets/audiocaps/prepare_audiocaps.py --out-dir data/audiocaps +``` + +The default manifests contain the full official train split and 64 distinct clips sampled +deterministically from the official validation split. Validation has five captions per clip; +the converter chooses one deterministically so a clip is not counted five times. Pass +`--keep-all-eval-captions --eval-limit 0` to retain every validation caption. + +For the event-grounded recipe, download `class_labels_indices.csv`, `eval_segments.csv`, +`balanced_train_segments.csv`, and `unbalanced_train_segments.csv` from the official +[AudioSet download page](https://research.google.com/audioset/download.html), place them in +one directory, and run: + +```bash +python datasets/audiocaps/prepare_audiocaps.py \ + --out-dir data/audiocaps-event-grounded \ + --audioset-metadata-dir /path/to/audioset-metadata +``` + +The converter requires an exact AudioSet segment match for every selected AudioCaps clip and +fails instead of silently emitting an ungrounded row. + +Each row uses the original AudioCaps caption as both the generation prompt and CLAP target: + +```json +{ + "prompt": "Multiple clanging and clanking sounds", + "metadata": { + "negative_audio_captions": ["... seven captions from other rows ..."], + "audio_event_labels": ["Door", "Sliding door"], + "audiocap_id": "58146", + "source_dataset": "AudioCaps", + "source_split": "train" + } +} +``` + +The CLAP scorer reads the positive text directly from `prompt`; no separate +audio-caption field or hand-written rewrite is required for this dataset. + +The deterministic negative captions provide the hardest-negative retrieval margin used by +the CLAP recipe and its retrieval diagnostics. The event-grounded recipe additionally scores +the least-aligned positive AudioSet label, encouraging every labeled event to remain audible. + +## Train + +```bash +DATA_PATH=data/audiocaps/train.jsonl \ +EVAL_DATA_PATH=data/audiocaps/eval.jsonl \ +bash examples/run_experiment_single_node.sh \ + diffusion/ltx2/ltx2_3_t2av_clap_trainside +``` + +Train and evaluation use the official AudioCaps train/validation split boundary. + +To train with RMS normalization and AudioSet event coverage: + +```bash +DATA_PATH=data/audiocaps-event-grounded/train.jsonl \ +EVAL_DATA_PATH=data/audiocaps-event-grounded/eval.jsonl \ +bash examples/run_experiment_single_node.sh \ + diffusion/ltx2/ltx2_3_t2av_clap_event_grounded_trainside +``` diff --git a/datasets/audiocaps/prepare_audiocaps.py b/datasets/audiocaps/prepare_audiocaps.py new file mode 100644 index 000000000..9dff02e0b --- /dev/null +++ b/datasets/audiocaps/prepare_audiocaps.py @@ -0,0 +1,197 @@ +"""Convert official AudioCaps captions to UniRL prompt JSONL files.""" + +from __future__ import annotations + +import argparse +import csv +import io +import json +import os +import random +import urllib.request +from typing import Dict, Iterable, List, Optional, Tuple + +DEFAULT_SOURCE = "https://raw.githubusercontent.com/cdjkim/audiocaps/d004db3ea1b01cf4fd0347dd8d27db90cadc8809" +AudioSetKey = Tuple[str, int] + + +def _audioset_key(youtube_id: str, start_time: str | float) -> AudioSetKey: + return str(youtube_id).strip(), round(float(start_time) * 1000) + + +def _read_audioset_labels(metadata_dir: str) -> Dict[AudioSetKey, List[str]]: + """Read the three official AudioSet segment files and resolve their label names.""" + label_path = os.path.join(metadata_dir, "class_labels_indices.csv") + with open(label_path, encoding="utf-8", newline="") as handle: + display_names = {row["mid"]: row["display_name"] for row in csv.DictReader(handle)} + + labels_by_segment: Dict[AudioSetKey, List[str]] = {} + filenames = ("eval_segments.csv", "balanced_train_segments.csv", "unbalanced_train_segments.csv") + for filename in filenames: + with open(os.path.join(metadata_dir, filename), encoding="utf-8", newline="") as handle: + rows = csv.reader((line for line in handle if not line.startswith("#")), skipinitialspace=True) + for row in rows: + key = _audioset_key(row[0], row[1]) + mids = [mid.strip() for mid in row[3].strip().strip('"').split(",")] + labels = [display_names[mid] for mid in mids] + previous = labels_by_segment.setdefault(key, labels) + if previous != labels: + raise ValueError(f"Conflicting AudioSet labels for segment {key!r}.") + return labels_by_segment + + +def _read_rows(source: str, split: str) -> List[Dict[str, str]]: + filename = f"{split}.csv" + if os.path.isdir(source): + with open(os.path.join(source, filename), encoding="utf-8", newline="") as handle: + return list(csv.DictReader(handle)) + + url = f"{source.rstrip('/')}/dataset/{filename}" + with urllib.request.urlopen(url, timeout=60) as response: # noqa: S310 - user-selectable dataset source + content = response.read().decode("utf-8-sig") + return list(csv.DictReader(io.StringIO(content))) + + +def _select_rows( + rows: List[Dict[str, str]], + limit: int, + seed: int, + *, + one_caption_per_clip: bool, +) -> List[Dict[str, str]]: + valid = [row for row in rows if str(row.get("caption") or "").strip()] + if one_caption_per_clip: + rows_by_clip: Dict[tuple[str, str], List[Dict[str, str]]] = {} + for row in valid: + rows_by_clip.setdefault((row["youtube_id"], row["start_time"]), []).append(row) + rng = random.Random(seed) + valid = [rng.choice(rows_by_clip[clip_key]) for clip_key in sorted(rows_by_clip)] + if limit <= 0 or limit >= len(valid): + return valid + indices = list(range(len(valid))) + random.Random(seed).shuffle(indices) + return [valid[index] for index in sorted(indices[:limit])] + + +def _negative_captions( + rows: List[Dict[str, str]], + index: int, + count: int, + *, + seed: int, + split: str, +) -> List[str]: + target = rows[index]["caption"].strip() + rng = random.Random(f"{seed}:{split}:{rows[index]['audiocap_id']}") + negatives: List[str] = [] + attempted_indices = set() + while len(attempted_indices) < len(rows): + candidate_index = rng.randrange(len(rows)) + if candidate_index in attempted_indices: + continue + attempted_indices.add(candidate_index) + caption = rows[candidate_index]["caption"].strip() + if candidate_index != index and caption != target and caption not in negatives: + negatives.append(caption) + if len(negatives) == count: + return negatives + raise ValueError(f"Could not find {count} distinct negative captions for AudioCaps row {index}.") + + +def _write_split( + rows: List[Dict[str, str]], + out_path: str, + *, + source_split: str, + negative_count: int, + seed: int, + audioset_labels: Optional[Dict[AudioSetKey, List[str]]], +) -> None: + os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True) + with open(out_path, "w", encoding="utf-8") as output: + for index, row in enumerate(rows): + caption = row["caption"].strip() + metadata = { + "negative_audio_captions": _negative_captions( + rows, + index, + negative_count, + seed=seed, + split=source_split, + ), + "audiocap_id": str(row["audiocap_id"]), + "youtube_id": str(row["youtube_id"]), + "start_time": float(row["start_time"]), + "source_dataset": "AudioCaps", + "source_split": source_split, + } + if audioset_labels is not None: + key = _audioset_key(row["youtube_id"], row["start_time"]) + if key not in audioset_labels: + raise ValueError(f"AudioCaps row {row['audiocap_id']} has no matching AudioSet segment {key!r}.") + metadata["audio_event_labels"] = audioset_labels[key] + record = { + "prompt": caption, + "prompt_id": f"audiocaps:{source_split}:{row['audiocap_id']}", + "metadata": metadata, + } + output.write(json.dumps(record, ensure_ascii=False) + "\n") + print(f"wrote {len(rows)} records -> {out_path}") + + +def _iter_specs(args: argparse.Namespace) -> Iterable[tuple[str, str, int, int, bool]]: + yield args.train_split, "train.jsonl", args.train_limit, args.seed, False + yield args.eval_split, "eval.jsonl", args.eval_limit, args.seed + 1, not args.keep_all_eval_captions + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", default=DEFAULT_SOURCE, help="AudioCaps repo URL or local directory") + parser.add_argument("--out-dir", default="data/audiocaps") + parser.add_argument("--train-split", default="train") + parser.add_argument("--eval-split", default="val") + parser.add_argument("--train-limit", type=int, default=0, help="0 keeps the full source split") + parser.add_argument("--eval-limit", type=int, default=64, help="0 keeps the full source split") + parser.add_argument("--negative-count", type=int, default=7) + parser.add_argument( + "--audioset-metadata-dir", + help="directory containing the three official AudioSet segment CSVs and class_labels_indices.csv", + ) + parser.add_argument( + "--keep-all-eval-captions", + action="store_true", + help="keep all five validation captions per clip instead of choosing one deterministically", + ) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if args.negative_count < 1: + parser.error("--negative-count must be at least 1") + + audioset_labels = _read_audioset_labels(args.audioset_metadata_dir) if args.audioset_metadata_dir else None + + for source_split, filename, limit, seed, one_caption_per_clip in _iter_specs(args): + source_rows = _read_rows(args.source, source_split) + selected_rows = _select_rows( + source_rows, + limit, + seed, + one_caption_per_clip=one_caption_per_clip, + ) + if len(selected_rows) <= args.negative_count: + parser.error( + f"{source_split} produced {len(selected_rows)} rows, but --negative-count={args.negative_count} " + "requires at least one more row" + ) + _write_split( + selected_rows, + os.path.join(args.out_dir, filename), + source_split=source_split, + negative_count=args.negative_count, + seed=args.seed, + audioset_labels=audioset_labels, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusion/ltx2/ltx2_3_t2av_clap_ast_event_grounded_trainside.yaml b/examples/diffusion/ltx2/ltx2_3_t2av_clap_ast_event_grounded_trainside.yaml new file mode 100644 index 000000000..e79f61dd3 --- /dev/null +++ b/examples/diffusion/ltx2/ltx2_3_t2av_clap_ast_event_grounded_trainside.yaml @@ -0,0 +1,19 @@ +# @package _global_ +# LTX-Video-2.3 T2AV Flow-GRPO with CLAP alignment plus an independent +# AudioSet AST event classifier. The AST term closes the gap where CLAP score +# rises while a concrete event such as applause becomes less audible. + +defaults: + - ltx2_3_t2av_clap_event_grounded_trainside + - _self_ + +save_dir: ${oc.env:SAVE_DIR,checkpoints/ltx2_3_t2av_clap_ast_event_grounded_trainside} + +logging: + tags: [ltx2.3, t2av, flowgrpo, trainside, audio, clap, ast, event-grounded, rms-normalized] + +reward: + backend: + config: + ast_event_weight: 1.0 + ast_model_id: MIT/ast-finetuned-audioset-10-10-0.4593 diff --git a/examples/diffusion/ltx2/ltx2_3_t2av_clap_event_grounded_trainside.yaml b/examples/diffusion/ltx2/ltx2_3_t2av_clap_event_grounded_trainside.yaml new file mode 100644 index 000000000..a8e06e35d --- /dev/null +++ b/examples/diffusion/ltx2/ltx2_3_t2av_clap_event_grounded_trainside.yaml @@ -0,0 +1,31 @@ +# @package _global_ +# LTX-Video-2.3 T2AV Flow-GRPO with loudness-invariant, event-grounded CLAP reward. +# +# Cook AudioCaps with official AudioSet segment metadata before launching: +# python datasets/audiocaps/prepare_audiocaps.py \ +# --out-dir data/audiocaps-event-grounded \ +# --audioset-metadata-dir /path/to/audioset-metadata +# +# Launch (1 node x 8 GPUs): +# DATA_PATH=data/audiocaps-event-grounded/train.jsonl \ +# EVAL_DATA_PATH=data/audiocaps-event-grounded/eval.jsonl \ +# bash examples/run_experiment_single_node.sh \ +# diffusion/ltx2/ltx2_3_t2av_clap_event_grounded_trainside + +defaults: + - ltx2_3_t2av_clap_trainside + - _self_ + +save_dir: ${oc.env:SAVE_DIR,checkpoints/ltx2_3_t2av_clap_event_grounded_trainside} + +logging: + tags: [ltx2.3, t2av, flowgrpo, trainside, audio, clap, event-grounded, rms-normalized] + +reward: + backend: + config: + event_prompts_metadata_key: audio_event_labels + event_coverage_weight: 1.0 + audio_normalization: rms + target_rms_dbfs: -20.0 + peak_limit: 0.95 diff --git a/examples/diffusion/ltx2/ltx2_3_t2av_clap_trainside.yaml b/examples/diffusion/ltx2/ltx2_3_t2av_clap_trainside.yaml new file mode 100644 index 000000000..6294cb8d1 --- /dev/null +++ b/examples/diffusion/ltx2/ltx2_3_t2av_clap_trainside.yaml @@ -0,0 +1,185 @@ +# @package _global_ +# LTX-Video-2.3 T2AV Flow-GRPO with audio-text reward only. +# +# The generated video remains the primary rollout output, while CLAP reads the +# decoded audio waveform carried alongside it in the reward request. Prepare the +# official AudioCaps captions first; each caption is both the generation prompt +# and CLAP target. Per-record distractor captions provide retrieval diagnostics +# without introducing a hand-written class vocabulary. +# +# Prepare data: +# python datasets/audiocaps/prepare_audiocaps.py --out-dir data/audiocaps +# +# Launch (1 node x 8 GPUs): +# DATA_PATH=data/audiocaps/train.jsonl EVAL_DATA_PATH=data/audiocaps/eval.jsonl \ +# bash examples/run_experiment_single_node.sh diffusion/ltx2/ltx2_3_t2av_clap_trainside + +num_devices: 8 +batch_size: 8 +adv_use_global_std: false +num_rollouts: 100 +eval_interval: 20 +eval_num_prompts: 8 +eval_samples_per_prompt: 1 +eval_chunk_prompts: 8 +eval_eta: 0.0 +eval_sampling: + init_same_noise: false + num_inference_steps: 30 + guidance_scale: 3.0 + height: 512 + width: 768 +save_interval: 20 +save_dir: ${oc.env:SAVE_DIR,checkpoints/ltx2_3_t2av_clap_trainside} +save_mode: adapter + +logging: + report_to_wandb: true + project_name: ${oc.env:WANDB_PROJECT,unirl-ltx2.3-t2av-clap} + run_name: null + entity: ${oc.env:WANDB_ENTITY,null} + tags: [ltx2.3, t2av, flowgrpo, trainside, audio, clap] + log_media: true + media_max_items: 4 + media_log_interval: 500 + +bundle: + _target_: unirl.models.ltx2.bundle.LTX2Bundle.from_config + config: + _target_: unirl.models.ltx2.config.LTX2PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL,dg845/LTX-2.3-Diffusers} + model_precision: bf16 + autocast_precision: bf16 + trajectory_precision: fp16 + logprob_precision: fp32 + shift: 1.0 + max_sequence_length: 512 + enable_audio: true + audio_joint_sde: true + audio_policy_logp_weight: 1.0 + default_height: 512 + default_width: 768 + default_num_frames: 49 + default_frame_rate: 24.0 + +pipeline: + _target_: unirl.models.ltx2.pipeline.LTX2Pipeline.from_bundle + config: ${bundle.config} + strategy: + _target_: unirl.sde.kernels.FlowSDEStrategy + +backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["LTX2VideoTransformerBlock"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 1.0e-4 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: constant + warmup_steps: 0 + total_steps: 10000 + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 32 + alpha: 256 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + target_modules: + - audio_proj_in + - audio_proj_out + - audio_attn1.to_q + - audio_attn1.to_k + - audio_attn1.to_v + - audio_attn1.to_out.0 + - audio_attn2.to_q + - audio_attn2.to_k + - audio_attn2.to_v + - audio_attn2.to_out.0 + - video_to_audio_attn.to_q + - video_to_audio_attn.to_k + - video_to_audio_attn.to_v + - video_to_audio_attn.to_out.0 + - audio_ff.net.0.proj + - audio_ff.net.2 + +rollout: + _target_: unirl.rollout.engine.trainside.engine.TrainsideRolloutEngine + stage_attrs: [diffusion] + forward_batch_size: 1 + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: unirl.reward.local.clap.CLAPRewardScorer + base_device: cuda + config: + _target_: unirl.reward.local.clap.CLAPSpec + batch_size: 2 + device: auto + negative_prompts_metadata_key: negative_audio_captions + matched_cosine_weight: 1.0 + retrieval_margin_weight: 1.0 + +algorithm: + _target_: unirl.algorithms.flowgrpo.FlowGRPO + stage_attr: diffusion + clip_range: 5.0e-3 + clip_schedule: constant + old_logp_source: replay + conditions_cls: + _target_: hydra.utils.get_class + path: unirl.models.ltx2.conditions.LTX2Conditions + params: ${sampling} + +stack: + _target_: unirl.train.stack.TrainStack + micro_batch_size: 1 + max_grad_norm: 1.0 + num_updates_per_batch: 2 + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH,data/audiocaps/train.jsonl} + eval_data_path: ${oc.env:EVAL_DATA_PATH,data/audiocaps/eval.jsonl} + seed: 42 + shuffle: true + algorithm: + prompts_per_rollout: ${batch_size} + +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + num_inference_steps: 10 + guidance_scale: 1.0 + height: 256 + width: 384 + num_frames: 49 + eta: 0.7 + samples_per_prompt: 8 + seed: 42 + init_same_noise: false + autocast_precision: bf16 + trajectory_precision: fp16 + logprob_precision: fp32 + scheduler: + _target_: unirl.sde.index_schedule.AllSDEScheduler + num_timesteps: ${..num_inference_steps} + num_sde_steps: 2 + timestep_fraction: [0, 0.5] diff --git a/examples/diffusion/ltx2/ltx2_3_t2av_verl_cps_clap_only_highbatch32_300.yaml b/examples/diffusion/ltx2/ltx2_3_t2av_verl_cps_clap_only_highbatch32_300.yaml new file mode 100644 index 000000000..0a094ee37 --- /dev/null +++ b/examples/diffusion/ltx2/ltx2_3_t2av_verl_cps_clap_only_highbatch32_300.yaml @@ -0,0 +1,188 @@ +# @package _global_ +# Fresh CLAP-only objective; micro-batch 4 retry after the micro-batch 8 OOM. +# 32 prompts x 8 samples = 256 samples/rollout; 2 updates of 128 samples each. +# Eight GPUs: micro-batch 4 x 4 accumulation micros per update on each GPU. +# Rollout forward batch stays 2; global optimizer batch remains 128. +# 150 rollouts = 300 optimizer updates; no baseline or periodic evaluation. +# Joint AV policy/LoRA and CPS parameters match the existing CLAP-only comparison. +# No advantage clipping; SDE steps use unirl's native AllSDEScheduler, restricted +# to the same first-10-of-24-step pool the old LTX2VerlIndexScheduler sampled +# from (3 random indices out of [0, 10)), reseeded per rollout step. + +num_devices: 8 +batch_size: 32 +adv_use_global_std: true +num_rollouts: 150 +load_dir: null +eval_interval: 0 +save_interval: 10 +save_dir: ${oc.env:LTX_CLAP_HIGHBATCH_SAVE_DIR,checkpoints/ltx23_verl_cps_clap_p32g8_u300_noeval} +save_mode: adapter +logging: + report_to_wandb: true + project_name: ${oc.env:WANDB_PROJECT,leo2_traing} + run_name: ${oc.env:WANDB_RUN_NAME,ltx23-clap-only-p32g8-mb4-n24-cfg4-cps08-allsde3-noclip-lr3e4-r64a128-f81-u300-noeval} + entity: ${oc.env:WANDB_ENTITY,boyeniu-the-university-of-sydney} + tags: + - ltx2.3 + - verl-reference + - cps + - clap-only + - highbatch32 + - joint-av + - allsde + - noclip + - noeval + log_media: true + media_max_items: 4 + media_log_interval: 10 +bundle: + _target_: unirl.models.ltx2.bundle.LTX2Bundle.from_config + config: + _target_: unirl.models.ltx2.config.LTX2PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL,diffusers/LTX-2.3-Diffusers} + model_precision: bf16 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + shift: 1.0 + max_sequence_length: 1024 + enable_audio: true + audio_joint_sde: true + audio_policy_logp_weight: null + default_height: ${sampling.height} + default_width: ${sampling.width} + default_num_frames: ${sampling.num_frames} + default_frame_rate: 24.0 +pipeline: + _target_: unirl.models.ltx2.pipeline.LTX2Pipeline.from_bundle + config: ${bundle.config} + strategy: + _target_: unirl.sde.kernels.CPSSDEStrategy +backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: + - LTX2VideoTransformerBlock + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + master_dtype: bf16 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 0.0003 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-08 + weight_decay: 0.0001 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: constant + warmup_steps: 0 + total_steps: 300 + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 64 + alpha: 128 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + target_modules: + - attn1.to_q + - attn1.to_k + - attn1.to_v + - attn1.to_out.0 + - attn2.to_q + - attn2.to_k + - attn2.to_v + - attn2.to_out.0 + - audio_attn1.to_q + - audio_attn1.to_k + - audio_attn1.to_v + - audio_attn1.to_out.0 + - audio_attn2.to_q + - audio_attn2.to_k + - audio_attn2.to_v + - audio_attn2.to_out.0 + - audio_to_video_attn.to_q + - audio_to_video_attn.to_k + - audio_to_video_attn.to_v + - audio_to_video_attn.to_out.0 + - video_to_audio_attn.to_q + - video_to_audio_attn.to_k + - video_to_audio_attn.to_v + - video_to_audio_attn.to_out.0 + - ff.net.0.proj + - ff.net.2 + - audio_ff.net.0.proj + - audio_ff.net.2 +rollout: + _target_: unirl.rollout.engine.trainside.engine.TrainsideRolloutEngine + stage_attrs: + - diffusion + forward_batch_size: 2 +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: unirl.reward.local.clap.CLAPRewardScorer + base_device: cuda + config: + _target_: unirl.reward.local.clap.CLAPSpec + batch_size: 4 + device: auto + model_id: laion/larger_clap_general + matched_cosine_weight: 1.0 + retrieval_margin_weight: 0.0 +algorithm: + _target_: unirl.algorithms.flowgrpo.FlowGRPO + stage_attr: diffusion + clip_range: 0.0001 + clip_schedule: constant + old_logp_source: replay + conditions_cls: + _target_: hydra.utils.get_class + path: unirl.models.ltx2.conditions.LTX2Conditions + params: ${sampling} +stack: + _target_: unirl.train.stack.TrainStack + micro_batch_size: 4 + max_grad_norm: 1.0 + num_updates_per_batch: 2 +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:LTX_VERL_DATA_PATH,data/vidprom_verl_reference/train.txt} + eval_data_path: null + seed: 42 + shuffle: true + algorithm: + prompts_per_rollout: ${batch_size} +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + num_inference_steps: 24 + guidance_scale: 4.0 + height: 256 + width: 384 + num_frames: 81 + eta: 0.8 + samples_per_prompt: 8 + seed: 42 + init_same_noise: false + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + scheduler: + _target_: unirl.sde.index_schedule.AllSDEScheduler + num_timesteps: ${..num_inference_steps} + # int(24 * 10/24) == 10, so this pins the pool to indices [0, 10) exactly, + # matching the old LTX2VerlIndexScheduler's window. + timestep_fraction: 0.4166666666666667 + num_sde_steps: 3 + max_sequence_length: 1024 diff --git a/unirl/models/ltx2/bundle.py b/unirl/models/ltx2/bundle.py index ae9ee7606..13ab09159 100644 --- a/unirl/models/ltx2/bundle.py +++ b/unirl/models/ltx2/bundle.py @@ -106,7 +106,7 @@ def from_config(cls, config: LTX2PipelineConfig) -> "LTX2Bundle": if config.enable_audio: try: from diffusers import AutoencoderKLLTX2Audio - from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder + from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder, LTX2VocoderWithBWE audio_vae = ( AutoencoderKLLTX2Audio.from_pretrained( @@ -117,8 +117,16 @@ def from_config(cls, config: LTX2PipelineConfig) -> "LTX2Bundle": ) audio_vae.requires_grad_(False) + vocoder_config = LTX2Vocoder.load_config(path, subfolder="vocoder") + vocoder_class_name = vocoder_config.get("_class_name", "LTX2Vocoder") + vocoder_cls = { + "LTX2Vocoder": LTX2Vocoder, + "LTX2VocoderWithBWE": LTX2VocoderWithBWE, + }.get(vocoder_class_name) + if vocoder_cls is None: + raise ValueError(f"LTX2Bundle: unsupported vocoder class {vocoder_class_name!r}.") vocoder = ( - LTX2Vocoder.from_pretrained(path, subfolder="vocoder", torch_dtype=dtype, low_cpu_mem_usage=False) + vocoder_cls.from_pretrained(path, subfolder="vocoder", torch_dtype=dtype, low_cpu_mem_usage=False) .to(device) .eval() ) diff --git a/unirl/models/ltx2/config.py b/unirl/models/ltx2/config.py index ded0f504a..79bc9ebee 100644 --- a/unirl/models/ltx2/config.py +++ b/unirl/models/ltx2/config.py @@ -36,6 +36,7 @@ class LTX2PipelineConfig: enable_audio: bool = False audio_joint_sde: bool = True + audio_policy_logp_weight: Optional[float] = None default_height: int = 512 default_width: int = 768 @@ -51,6 +52,13 @@ class LTX2PipelineConfig: def __post_init__(self) -> None: validate_precision_type(self.model_precision, field="LTX2PipelineConfig.model_precision") + if self.audio_policy_logp_weight is not None: + self.audio_policy_logp_weight = float(self.audio_policy_logp_weight) + if not 0.0 <= self.audio_policy_logp_weight <= 1.0: + raise ValueError( + "LTX2PipelineConfig.audio_policy_logp_weight must be in [0, 1], " + f"got {self.audio_policy_logp_weight}." + ) __all__ = ["LTX2PipelineConfig"] diff --git a/unirl/models/ltx2/diffusion.py b/unirl/models/ltx2/diffusion.py index 5a5e9e64b..60b40a574 100644 --- a/unirl/models/ltx2/diffusion.py +++ b/unirl/models/ltx2/diffusion.py @@ -74,10 +74,13 @@ def _combine_modality_logp( audio_logp: torch.Tensor, n_video: int, n_audio: int, + audio_weight: Optional[float] = None, ) -> torch.Tensor: - """Element-weighted mean of the per-step video/audio log-probs.""" - total = n_video + n_audio - return (video_logp * n_video + audio_logp * n_audio) / total + """Combine per-step video/audio mean log-probs using element or explicit modality weighting.""" + if audio_weight is None: + total = n_video + n_audio + return (video_logp * n_video + audio_logp * n_audio) / total + return video_logp * (1.0 - audio_weight) + audio_logp * audio_weight class LTX2DiffusionStep(DiffusionStep[LTX2Bundle, LTX2Conditions]): @@ -183,6 +186,7 @@ def __init__( trajectory_precision: str = "fp16", logprob_precision: str = "fp32", audio_joint_sde: bool = True, + audio_policy_logp_weight: Optional[float] = None, ) -> None: self.bundle = bundle self.step_kernel = LTX2DiffusionStep() @@ -191,6 +195,11 @@ def __init__( self.trajectory_dtype = parse_torch_dtype(trajectory_precision, field_name="trajectory_precision") self.logprob_dtype = parse_torch_dtype(logprob_precision, field_name="logprob_precision") self.audio_joint_sde = bool(audio_joint_sde) + if audio_policy_logp_weight is not None: + audio_policy_logp_weight = float(audio_policy_logp_weight) + if not 0.0 <= audio_policy_logp_weight <= 1.0: + raise ValueError(f"audio_policy_logp_weight must be in [0, 1], got {audio_policy_logp_weight}.") + self.audio_policy_logp_weight = audio_policy_logp_weight self._audio_in_policy = self.audio_joint_sde and bool(getattr(bundle, "has_audio", False)) def trainable_module(self) -> torch.nn.Module: @@ -274,6 +283,15 @@ def generate( if step_eta > 0.0 and denoise_seed_keys is not None else None ) + audio_step_generators = ( + make_denoise_step_generators( + base_seed=int(denoise_base_seed), + step_index=step_idx, + sample_ids=[f"{key}::audio" for key in denoise_seed_keys], + ) + if step_eta > 0.0 and self._audio_in_policy and denoise_seed_keys is not None + else None + ) video_pred, audio_pred = self.step_kernel.predict_noise( self.bundle, @@ -307,6 +325,7 @@ def generate( sigma=sigma, sigma_next=sigma_next, eta=audio_eta, + generator=audio_step_generators, sigma_max=sigma_max, step_index=step_idx, ) @@ -322,6 +341,7 @@ def generate( audio_log_prob, n_video=x[0].numel(), n_audio=a[0].numel(), + audio_weight=self.audio_policy_logp_weight, ) sde_logp_list.append(log_prob.to(dtype=self.logprob_dtype)) @@ -441,6 +461,7 @@ def replay( audio_log_prob, n_video=sample[0].numel(), n_audio=audio_sample[0].numel(), + audio_weight=self.audio_policy_logp_weight, ) if prev_mean is not None and audio_prev_mean is not None: prev_mean = torch.cat([prev_mean, audio_prev_mean], dim=1) diff --git a/unirl/models/ltx2/pipeline.py b/unirl/models/ltx2/pipeline.py index ebce8514a..80d641d3b 100644 --- a/unirl/models/ltx2/pipeline.py +++ b/unirl/models/ltx2/pipeline.py @@ -90,6 +90,7 @@ def from_bundle( trajectory_precision=config.trajectory_precision, logprob_precision=config.logprob_precision, audio_joint_sde=config.audio_joint_sde, + audio_policy_logp_weight=config.audio_policy_logp_weight, ) vae_decode = LTX2VAEDecodeStage(bundle) vae_encode = LTX2VAEEncodeStage(bundle) @@ -227,6 +228,7 @@ def generate(self, sample: Sample) -> Sample: ) sde_indices = list(params.sde_indices) if params.sde_indices is not None else None + denoise_seed_keys = list(frontier.init_noise_group_ids or frontier.sample_ids) segment = self.diffusion.generate( conditions, params=params, @@ -234,7 +236,7 @@ def generate(self, sample: Sample) -> Sample: initial_latents=initial_latents, initial_audio_latents=initial_audio_latents, sde_indices=sde_indices, - denoise_seed_keys=[str(sample_id) for sample_id in sample.sample_ids], + denoise_seed_keys=[str(seed_key) for seed_key in denoise_seed_keys], denoise_base_seed=int(params.seed) if params.seed is not None else 0, ) @@ -263,7 +265,7 @@ def generate(self, sample: Sample) -> Sample: audio_sample_rate = int(self.bundle.vocoder.config.output_sampling_rate) primitives = {"video": decoded} - primitive_metadata = {} + primitive_metadata = {"video": {"fps": float(self.config.default_frame_rate)}} if decoded_audio is not None: primitives["audio"] = decoded_audio primitive_metadata["audio"] = {"sample_rate": audio_sample_rate} diff --git a/unirl/reward/local/clap.py b/unirl/reward/local/clap.py index 84949ad8b..17e4b2f24 100644 --- a/unirl/reward/local/clap.py +++ b/unirl/reward/local/clap.py @@ -2,15 +2,18 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import List +import inspect +import math +import time +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional import torch import torch.nn.functional as F from unirl.reward.base import BaseRewardComponentSpec from unirl.reward.local.device import resolve_device -from unirl.types.reward import RewardRequest +from unirl.types.reward import RewardRequest, RewardResponse from .base import LocalRewardBackend @@ -23,6 +26,59 @@ class CLAPRewardScorer(LocalRewardBackend): CLAP_SAMPLE_RATE = 48_000 def __init__(self, *, config: "CLAPSpec", base_device: str) -> None: + self.prompt_metadata_key = str(config.prompt_metadata_key or "").strip() or None + self.negative_prompts_metadata_key = str(config.negative_prompts_metadata_key or "").strip() or None + self.event_prompts_metadata_key = str(config.event_prompts_metadata_key or "").strip() or None + self.matched_cosine_weight = float(config.matched_cosine_weight) + self.retrieval_margin_weight = float(config.retrieval_margin_weight) + self.event_coverage_weight = float(config.event_coverage_weight) + self.ast_event_weight = float(config.ast_event_weight) + self.ast_model_id = str(config.ast_model_id).strip() + reward_weights = ( + self.matched_cosine_weight, + self.retrieval_margin_weight, + self.event_coverage_weight, + self.ast_event_weight, + ) + if not all(math.isfinite(weight) for weight in reward_weights): + raise ValueError("CLAP reward weights must be finite.") + if any(weight < 0.0 for weight in reward_weights): + raise ValueError("CLAP reward weights must be non-negative.") + if not any(weight > 0.0 for weight in reward_weights): + raise ValueError("At least one CLAP reward weight must be positive.") + self.audio_normalization = str(config.audio_normalization).strip().lower() + self.target_rms_dbfs = float(config.target_rms_dbfs) + self.peak_limit = float(config.peak_limit) + if self.audio_normalization not in {"none", "rms"}: + raise ValueError("CLAPSpec.audio_normalization must be 'none' or 'rms'.") + if not math.isfinite(self.target_rms_dbfs) or self.target_rms_dbfs > 0.0: + raise ValueError("CLAPSpec.target_rms_dbfs must be finite and at most 0 dBFS.") + if not math.isfinite(self.peak_limit) or not 0.0 < self.peak_limit <= 1.0: + raise ValueError("CLAPSpec.peak_limit must be in (0, 1].") + self.retrieval_prompts = [str(prompt).strip() for prompt in config.retrieval_prompts] + if any(not prompt for prompt in self.retrieval_prompts): + raise ValueError("CLAPSpec.retrieval_prompts must contain only non-empty strings.") + if len(set(self.retrieval_prompts)) != len(self.retrieval_prompts): + raise ValueError("CLAPSpec.retrieval_prompts must not contain duplicates.") + if self.retrieval_prompts and len(self.retrieval_prompts) < 2: + raise ValueError("CLAPSpec.retrieval_prompts needs at least two prompts for retrieval diagnostics.") + if self.retrieval_prompts and self.negative_prompts_metadata_key: + raise ValueError("CLAPSpec.retrieval_prompts and negative_prompts_metadata_key are mutually exclusive.") + if self.retrieval_margin_weight > 0.0 and not (self.retrieval_prompts or self.negative_prompts_metadata_key): + raise ValueError( + "CLAPSpec.retrieval_margin_weight requires retrieval_prompts or negative_prompts_metadata_key." + ) + if self.event_coverage_weight > 0.0 and not self.event_prompts_metadata_key: + raise ValueError("CLAPSpec.event_coverage_weight requires event_prompts_metadata_key.") + if self.ast_event_weight > 0.0 and not self.event_prompts_metadata_key: + raise ValueError("CLAPSpec.ast_event_weight requires event_prompts_metadata_key.") + if self.ast_event_weight > 0.0 and not self.ast_model_id: + raise ValueError("CLAPSpec.ast_model_id must be non-empty when ast_event_weight is positive.") + self._retrieval_text_embeds: Optional[torch.Tensor] = None + self.ast_model = None + self.ast_processor = None + self._ast_sample_rate: Optional[int] = None + self._ast_label_to_index: Dict[str, int] = {} super().__init__( device=resolve_device(config.device, base_device), batch_size=config.batch_size, @@ -39,11 +95,33 @@ def _load_model(self) -> None: self.model = ClapModel.from_pretrained(model_id).to(self.device).eval() self.model = self.model.to(dtype=torch.float32) self.processor = ClapProcessor.from_pretrained(model_id) + self._processor_audio_keyword = ( + "audio" if "audio" in inspect.signature(self.processor.__call__).parameters else "audios" + ) + if self.ast_event_weight > 0.0: + try: + from transformers import AutoFeatureExtractor, AutoModelForAudioClassification + except ImportError as e: + raise ImportError( + "transformers with AutoModelForAudioClassification is required for the AST event reward" + ) from e + + self.ast_processor = AutoFeatureExtractor.from_pretrained(self.ast_model_id) + self.ast_model = ( + AutoModelForAudioClassification.from_pretrained(self.ast_model_id) + .to(self.device) + .eval() + .to(dtype=torch.float32) + ) + self._ast_sample_rate = int(self.ast_processor.sampling_rate) + self._ast_label_to_index = { + str(label).strip().lower(): int(index) + for index, label in self.ast_model.config.id2label.items() + } def _preprocess_audio(self, audio_list: List[torch.Tensor], src_sample_rate: int) -> List["torch.Tensor"]: """Downmix and resample each ``[L]`` / ``[C, L]`` / ``[L, C]`` waveform to CLAP's 48 kHz mono ``[L']``.""" import numpy as np - import torchaudio.functional as AF processed: List[np.ndarray] = [] for waveform in audio_list: @@ -56,18 +134,228 @@ def _preprocess_audio(self, audio_list: List[torch.Tensor], src_sample_rate: int wf = wf.reshape(-1) if src_sample_rate != self.CLAP_SAMPLE_RATE: + import torchaudio.functional as AF + wf = AF.resample( wf.unsqueeze(0), orig_freq=int(src_sample_rate), new_freq=self.CLAP_SAMPLE_RATE, ).squeeze(0) + if self.audio_normalization == "rms": + rms = wf.square().mean().sqrt() + peak = wf.abs().max() + if rms > torch.finfo(wf.dtype).eps: + gain = (10.0 ** (self.target_rms_dbfs / 20.0)) / rms + if peak * gain > self.peak_limit: + gain = self.peak_limit / peak + wf = wf * gain + processed.append(wf.cpu().numpy()) return processed + @staticmethod + def _feature_tensor(value: Any) -> torch.Tensor: + """Normalize the tensor/wrapper return types used across transformers releases.""" + if torch.is_tensor(value): + return value + pooled = getattr(value, "pooler_output", None) + if torch.is_tensor(pooled): + return pooled + if isinstance(value, tuple) and value and torch.is_tensor(value[0]): + return value[0] + raise TypeError(f"Cannot extract CLAP feature tensor from {type(value).__name__}.") + + def _encode_texts(self, prompts: List[str]) -> torch.Tensor: + inputs = self.processor(text=prompts, return_tensors="pt", padding=True) + inputs = {key: value.to(self.device) for key, value in inputs.items()} + with torch.no_grad(): + features = self.model.get_text_features( + input_ids=inputs.get("input_ids"), + attention_mask=inputs.get("attention_mask"), + ) + return F.normalize(self._feature_tensor(features).float(), p=2, dim=-1) + + def _encode_audio(self, waveforms: List[torch.Tensor], src_sample_rate: int) -> torch.Tensor: + waveforms_np = self._preprocess_audio(waveforms, src_sample_rate) + inputs = self.processor( + **{self._processor_audio_keyword: waveforms_np}, + sampling_rate=self.CLAP_SAMPLE_RATE, + return_tensors="pt", + padding=True, + ) + inputs = {key: value.to(self.device) for key, value in inputs.items()} + with torch.no_grad(): + features = self.model.get_audio_features( + input_features=inputs.get("input_features"), + is_longer=inputs.get("is_longer"), + attention_mask=inputs.get("attention_mask"), + ) + return F.normalize(self._feature_tensor(features).float(), p=2, dim=-1) + + def _encode_ast_events( + self, + waveforms: List[torch.Tensor], + src_sample_rate: int, + event_prompts: List[List[str]], + ) -> torch.Tensor: + """Return the minimum AudioSet event probability requested by each sample.""" + if self.ast_model is None or self.ast_processor is None or self._ast_sample_rate is None: + raise RuntimeError("AST event reward is enabled but the AST model is not loaded.") + + import numpy as np + processed: List[np.ndarray] = [] + for waveform in waveforms: + wf = waveform.detach().float() + if wf.isnan().any() or wf.isinf().any(): + wf = torch.zeros_like(wf) + if wf.ndim == 2: + channel_axis = 0 if wf.shape[0] <= wf.shape[1] else 1 + wf = wf.mean(dim=channel_axis) + wf = wf.reshape(-1) + if src_sample_rate != self._ast_sample_rate: + if src_sample_rate % self._ast_sample_rate == 0: + factor = src_sample_rate // self._ast_sample_rate + wf = F.avg_pool1d( + wf.reshape(1, 1, -1), + kernel_size=factor, + stride=factor, + ).reshape(-1) + else: + target_length = max(1, round(wf.numel() * self._ast_sample_rate / src_sample_rate)) + wf = F.interpolate( + wf.reshape(1, 1, -1), + size=target_length, + mode="linear", + align_corners=False, + ).reshape(-1) + if self.audio_normalization == "rms": + rms = wf.square().mean().sqrt() + peak = wf.abs().max() + if rms > torch.finfo(wf.dtype).eps: + gain = (10.0 ** (self.target_rms_dbfs / 20.0)) / rms + if peak * gain > self.peak_limit: + gain = self.peak_limit / peak + wf = wf * gain + processed.append(wf.cpu().numpy()) + + inputs = self.ast_processor( + processed, + sampling_rate=self._ast_sample_rate, + return_tensors="pt", + padding=True, + ) + inputs = {key: value.to(self.device) for key, value in inputs.items()} + with torch.no_grad(): + probabilities = self.ast_model(**inputs).logits.float().sigmoid() + + scores = [] + for row_index, labels in enumerate(event_prompts): + missing = [label for label in labels if label.strip().lower() not in self._ast_label_to_index] + if missing: + raise ValueError(f"AST model does not define AudioSet event labels: {missing!r}.") + indices = torch.tensor( + [self._ast_label_to_index[label.strip().lower()] for label in labels], + device=probabilities.device, + dtype=torch.long, + ) + scores.append(probabilities[row_index].index_select(0, indices).min()) + return torch.stack(scores) + + def _reward_prompts(self, request: RewardRequest) -> List[str]: + prompts = request.prompts + if self.prompt_metadata_key is None: + return prompts + + metadata = request.metadata or [] + resolved: List[str] = [] + for index, prompt in enumerate(prompts): + row = metadata[index] if index < len(metadata) else None + candidate = row.get(self.prompt_metadata_key) if isinstance(row, dict) else None + resolved.append(candidate.strip() if isinstance(candidate, str) and candidate.strip() else prompt) + return resolved + + def _negative_prompts(self, request: RewardRequest, reward_prompts: List[str]) -> Optional[List[List[str]]]: + if self.negative_prompts_metadata_key is None: + return None + + metadata = request.metadata or [] + resolved: List[List[str]] = [] + for index, target_prompt in enumerate(reward_prompts): + row = metadata[index] if index < len(metadata) else None + candidates = row.get(self.negative_prompts_metadata_key) if isinstance(row, dict) else None + if not isinstance(candidates, (list, tuple)): + raise ValueError( + "CLAP negative prompt metadata must be a list of strings; " + f"key={self.negative_prompts_metadata_key!r}, sample_index={index}." + ) + negatives = list( + dict.fromkeys( + candidate.strip() + for candidate in candidates + if isinstance(candidate, str) and candidate.strip() and candidate.strip() != target_prompt + ) + ) + if not negatives: + raise ValueError( + "CLAP negative prompt metadata must contain at least one non-target caption; " + f"key={self.negative_prompts_metadata_key!r}, sample_index={index}." + ) + resolved.append(negatives) + return resolved + + def _event_prompts(self, request: RewardRequest) -> Optional[List[List[str]]]: + """Resolve per-sample sound-event prompts used for minimum-coverage scoring.""" + if self.event_prompts_metadata_key is None: + return None + + metadata = request.metadata or [] + resolved: List[List[str]] = [] + for index in range(len(request.prompts)): + row = metadata[index] if index < len(metadata) else None + candidates = row.get(self.event_prompts_metadata_key) if isinstance(row, dict) else None + if not isinstance(candidates, (list, tuple)): + raise ValueError( + "CLAP event prompt metadata must be a list of strings; " + f"key={self.event_prompts_metadata_key!r}, sample_index={index}." + ) + event_prompts = list( + dict.fromkeys( + candidate.strip() for candidate in candidates if isinstance(candidate, str) and candidate.strip() + ) + ) + if not event_prompts: + raise ValueError( + "CLAP event prompt metadata must contain at least one event; " + f"key={self.event_prompts_metadata_key!r}, sample_index={index}." + ) + resolved.append(event_prompts) + return resolved + + def compute_rewards(self, request: RewardRequest) -> RewardResponse: + """Score matched captions and expose optional retrieval diagnostics.""" + if not self._is_loaded: + raise RuntimeError( + f"{type(self).__name__}.compute_rewards called before _load_model " + f"completed (model_name={self.model_name!r}, batch_size={request.batch_size})." + ) + start = time.time() + rewards, components = self._compute_rewards_and_components(request) + return RewardResponse( + rewards=rewards, + component_rewards=components, + successes=[True] * len(rewards), + errors=[None] * len(rewards), + compute_time=time.time() - start, + ) + def _compute_model_rewards(self, request: RewardRequest) -> List[float]: + rewards, _ = self._compute_rewards_and_components(request) + return rewards + + def _compute_rewards_and_components(self, request: RewardRequest) -> tuple[List[float], Dict[str, List[float]]]: audio = request.audio - prompts = request.prompts + prompts = self._reward_prompts(request) if audio is None: raise ValueError( "CLAPRewardScorer requires audio in the reward request " @@ -76,35 +364,166 @@ def _compute_model_rewards(self, request: RewardRequest) -> List[float]: ) if request.audio_sample_rate is None: raise ValueError("CLAPRewardScorer requires request.audio_sample_rate (source Hz); got None.") + if len(audio) != len(prompts): + raise ValueError(f"CLAPRewardScorer got {len(audio)} audio samples but {len(prompts)} reward prompts.") src_rate = int(request.audio_sample_rate) + negative_prompt_sets = self._negative_prompts(request, prompts) + event_prompt_sets = self._event_prompts(request) + dynamic_text_embeds: Optional[torch.Tensor] = None + dynamic_prompt_to_index: Dict[str, int] = {} + if not self.retrieval_prompts or event_prompt_sets is not None: + dynamic_prompts = [] if self.retrieval_prompts else list(prompts) + if negative_prompt_sets is not None: + dynamic_prompts.extend(prompt for row in negative_prompt_sets for prompt in row) + if event_prompt_sets is not None: + dynamic_prompts.extend(prompt for row in event_prompt_sets for prompt in row) + unique_dynamic_prompts = list(dict.fromkeys(dynamic_prompts)) + dynamic_text_embeds = self._encode_texts(unique_dynamic_prompts) + dynamic_prompt_to_index = {prompt: index for index, prompt in enumerate(unique_dynamic_prompts)} all_rewards: List[float] = [] + component_rewards: Dict[str, List[float]] = {"matched_cosine": []} + retrieval_prompt_to_index = {prompt: index for index, prompt in enumerate(self.retrieval_prompts)} + if self.retrieval_prompts or negative_prompt_sets is not None: + component_rewards.update( + { + "mismatched_cosine": [], + "retrieval_margin": [], + "retrieval_top1": [], + } + ) + if event_prompt_sets is not None: + component_rewards["event_coverage_min_cosine"] = [] + if self.ast_event_weight > 0.0: + component_rewards["ast_event_min_probability"] = [] + if self.retrieval_prompts: + unknown = sorted(set(prompts) - set(retrieval_prompt_to_index)) + if unknown: + raise ValueError( + "CLAP reward prompts must appear in CLAPSpec.retrieval_prompts when retrieval diagnostics " + f"are enabled; missing={unknown[:3]!r}." + ) + for i in range(0, len(audio), self.batch_size): batch_audio = audio[i : i + self.batch_size] batch_prompts = prompts[i : i + self.batch_size] - waveforms_np = self._preprocess_audio(batch_audio, src_rate) - - inputs = self.processor( - text=batch_prompts, - # `audios=` was deprecated and is now rejected outright by - # ClapProcessor (transformers 5.x): "You passed keyword argument - # `audios` which is deprecated. Please use `audio` instead." - # It surfaces as every sample failing scoring, not as a crash. - audio=waveforms_np, - sampling_rate=self.CLAP_SAMPLE_RATE, - return_tensors="pt", - padding=True, - ) - inputs = {k: v.to(self.device) for k, v in inputs.items()} + audio_embeds = self._encode_audio(batch_audio, src_rate) + + if self.retrieval_prompts: + score_matrix = audio_embeds @ self._get_retrieval_text_embeds().T + target_indices = torch.tensor( + [retrieval_prompt_to_index[prompt] for prompt in batch_prompts], + device=score_matrix.device, + dtype=torch.long, + ) + rows = torch.arange(score_matrix.shape[0], device=score_matrix.device) + matched = score_matrix[rows, target_indices] + negative_mask = torch.ones_like(score_matrix, dtype=torch.bool) + negative_mask[rows, target_indices] = False + negative_scores = score_matrix[negative_mask].view(score_matrix.shape[0], -1) + mismatched = negative_scores.mean(dim=-1) + margin = matched - negative_scores.max(dim=-1).values + top1 = (score_matrix.argmax(dim=-1) == target_indices).float() + elif negative_prompt_sets is not None: + assert dynamic_text_embeds is not None + target_indices = torch.tensor( + [dynamic_prompt_to_index[prompt] for prompt in batch_prompts], + device=dynamic_text_embeds.device, + dtype=torch.long, + ) + text_embeds = dynamic_text_embeds.index_select(0, target_indices) + matched = (audio_embeds * text_embeds).sum(dim=-1) + batch_negative_prompts = negative_prompt_sets[i : i + self.batch_size] + mismatched_values = [] + hardest_negative_values = [] + for row_index, row_prompts in enumerate(batch_negative_prompts): + indices = torch.tensor( + [dynamic_prompt_to_index[prompt] for prompt in row_prompts], + device=dynamic_text_embeds.device, + dtype=torch.long, + ) + row_negative_embeds = dynamic_text_embeds.index_select(0, indices) + row_negative_scores = row_negative_embeds @ audio_embeds[row_index] + mismatched_values.append(row_negative_scores.mean()) + hardest_negative_values.append(row_negative_scores.max()) + mismatched = torch.stack(mismatched_values) + hardest_negative = torch.stack(hardest_negative_values) + margin = matched - hardest_negative + top1 = (matched > hardest_negative).float() + else: + assert dynamic_text_embeds is not None + target_indices = torch.tensor( + [dynamic_prompt_to_index[prompt] for prompt in batch_prompts], + device=dynamic_text_embeds.device, + dtype=torch.long, + ) + text_embeds = dynamic_text_embeds.index_select(0, target_indices) + matched = (audio_embeds * text_embeds).sum(dim=-1) + + event_coverage: Optional[torch.Tensor] = None + if event_prompt_sets is not None: + assert dynamic_text_embeds is not None + event_values = [] + for row_index, row_prompts in enumerate(event_prompt_sets[i : i + self.batch_size]): + indices = torch.tensor( + [dynamic_prompt_to_index[prompt] for prompt in row_prompts], + device=dynamic_text_embeds.device, + dtype=torch.long, + ) + row_event_embeds = dynamic_text_embeds.index_select(0, indices) + event_values.append((row_event_embeds @ audio_embeds[row_index]).min()) + event_coverage = torch.stack(event_values) + component_rewards["event_coverage_min_cosine"].extend(event_coverage.float().cpu().tolist()) + + ast_event_score: Optional[torch.Tensor] = None + if self.ast_event_weight > 0.0: + assert event_prompt_sets is not None + ast_event_score = self._encode_ast_events( + batch_audio, + src_rate, + event_prompt_sets[i : i + self.batch_size], + ) + component_rewards["ast_event_min_probability"].extend(ast_event_score.float().cpu().tolist()) + + if self.retrieval_prompts or negative_prompt_sets is not None: + component_rewards["mismatched_cosine"].extend(mismatched.float().cpu().tolist()) + component_rewards["retrieval_margin"].extend(margin.float().cpu().tolist()) + component_rewards["retrieval_top1"].extend(top1.cpu().tolist()) + + reward = matched * self.matched_cosine_weight + if self.retrieval_margin_weight > 0.0: + reward = reward + margin * self.retrieval_margin_weight + if self.event_coverage_weight > 0.0: + assert event_coverage is not None + reward = reward + event_coverage * self.event_coverage_weight + if self.ast_event_weight > 0.0: + assert ast_event_score is not None + reward = reward + ast_event_score * self.ast_event_weight + + matched_values = matched.float().cpu().tolist() + all_rewards.extend(reward.float().cpu().tolist()) + component_rewards["matched_cosine"].extend(matched_values) + + return all_rewards, component_rewards + + def _get_retrieval_text_embeds(self) -> torch.Tensor: + if self._retrieval_text_embeds is None: + self._retrieval_text_embeds = self._encode_texts(self.retrieval_prompts) + return self._retrieval_text_embeds - with torch.no_grad(): - outputs = self.model(**inputs) - audio_embeds = F.normalize(outputs.audio_embeds, p=2, dim=-1) - text_embeds = F.normalize(outputs.text_embeds, p=2, dim=-1) - scores = (audio_embeds * text_embeds).sum(dim=-1) - all_rewards.extend(scores.float().cpu().tolist()) + def offload(self) -> None: + super().offload() + if self.ast_model is not None: + self.ast_model = self.ast_model.cpu() + if self._retrieval_text_embeds is not None: + self._retrieval_text_embeds = self._retrieval_text_embeds.cpu() - return all_rewards + def onload(self) -> None: + super().onload() + if self.ast_model is not None: + self.ast_model = self.ast_model.to(self.device) + if self._retrieval_text_embeds is not None: + self._retrieval_text_embeds = self._retrieval_text_embeds.to(self.device) @dataclass @@ -114,3 +533,15 @@ class CLAPSpec(BaseRewardComponentSpec): batch_size: int = 8 device: str = "auto" model_id: str = "laion/larger_clap_general" + prompt_metadata_key: Optional[str] = None + negative_prompts_metadata_key: Optional[str] = None + event_prompts_metadata_key: Optional[str] = None + retrieval_prompts: List[str] = field(default_factory=list) + matched_cosine_weight: float = 1.0 + retrieval_margin_weight: float = 0.0 + event_coverage_weight: float = 0.0 + ast_event_weight: float = 0.0 + ast_model_id: str = "MIT/ast-finetuned-audioset-10-10-0.4593" + audio_normalization: str = "none" + target_rms_dbfs: float = -20.0 + peak_limit: float = 0.95 diff --git a/unirl/types/media_preview.py b/unirl/types/media_preview.py index b30263868..10c41ef12 100644 --- a/unirl/types/media_preview.py +++ b/unirl/types/media_preview.py @@ -21,6 +21,7 @@ class MediaPreview(Batch): images: List[Any] = concat_field(default_factory=list) videos: List[Any] = concat_field(default_factory=list) audios: List[Any] = concat_field(default_factory=list) + video_fps: Optional[float] = None audio_sample_rate: Optional[int] = None prompts: List[str] = concat_field(default_factory=list) rewards: List[float] = concat_field(default_factory=list) @@ -157,7 +158,14 @@ def build_media_preview_for_part( return None audios_out: List[Any] = [] + video_fps: Optional[float] = None audio_sr: Optional[int] = None + if videos: + raw_video_fps = part.primitive_metadata.get("video", {}).get("fps") + if raw_video_fps is not None: + video_fps = float(raw_video_fps) + if video_fps <= 0.0: + raise ValueError(f"Video preview fps must be > 0, got {video_fps}") decoded_audio = part.primitives.get("audio") if isinstance(decoded_audio, Audios): from unirl.distributed.tensor import hydrate, map_tree @@ -180,6 +188,7 @@ def build_media_preview_for_part( images=images, videos=videos, audios=audios_out, + video_fps=video_fps, audio_sample_rate=int(audio_sr) if audio_sr is not None else None, prompts=prompts_out, rewards=reward_values, diff --git a/unirl/utils/wandb_logger.py b/unirl/utils/wandb_logger.py index e1c4cde68..e19b534e0 100644 --- a/unirl/utils/wandb_logger.py +++ b/unirl/utils/wandb_logger.py @@ -406,7 +406,7 @@ def log_generated_media( *, key: str = "rollout/generated_media", video_key: Optional[str] = None, - video_fps: int = 8, + video_fps: Optional[float] = None, step_key: str = "rollout/step", ) -> None: """Log rollout media preview payload produced by the rollout pipeline.""" @@ -424,6 +424,17 @@ def log_generated_media( prompts = getattr(media_preview, "prompts", None) rewards = getattr(media_preview, "rewards", None) + preview_video_fps = ( + media_preview.get("video_fps") + if isinstance(media_preview, dict) + else getattr(media_preview, "video_fps", None) + ) + resolved_video_fps = float( + video_fps if video_fps is not None else (preview_video_fps if preview_video_fps is not None else 8) + ) + if resolved_video_fps <= 0.0: + raise ValueError(f"log_generated_media: video_fps must be > 0, got {resolved_video_fps}") + has_images = isinstance(images, list) and bool(images) has_videos = isinstance(videos, list) and bool(videos) if not has_images and not has_videos: @@ -501,11 +512,13 @@ def _caption_for(idx: int) -> str: audio_wf = audios[idx] if idx < len(audios) else None if audio_wf is not None and audio_sr is not None and torch.is_tensor(audio_wf): arr_hwc = arr.transpose(0, 2, 3, 1) # (T, C, H, W) -> (T, H, W, C) - path = _write_video_with_audio(arr_hwc, int(video_fps), audio_wf, int(audio_sr)) + path = _write_video_with_audio(arr_hwc, int(round(resolved_video_fps)), audio_wf, int(audio_sr)) _muxed_paths.append(path) wandb_videos.append(wandb.Video(path, caption=_caption_for(idx), format="mp4")) else: - wandb_videos.append(wandb.Video(arr, caption=_caption_for(idx), fps=int(video_fps))) + wandb_videos.append( + wandb.Video(arr, caption=_caption_for(idx), fps=int(round(resolved_video_fps))) + ) if wandb_videos: payload[video_key] = wandb_videos